Discovering Data Structures: Nearest Neighbor Search and Beyond

## Abstract

Abstract We propose a general framework for end-to-end learning of data structures. Our framework adapts to the underlying data distribution and provides fine-grained control over query and space complexity. Crucially, the data structure is learned from scratch, and does not require careful initialization or seeding with candidate data structures.
We first apply this framework to the problem of nearest neighbor search. In several settings, we are able to reverse-engineer the learned data structures and query algorithms. For 1D nearest neighbor search, the model discovers optimal distribution (in)dependent algorithms such as binary search and variants of interpolation search. In higher dimensions, the model learns solutions that resemble k-d trees in some regimes, while in others, elements of locality-sensitive hashing emerge. Additionally, the model learns useful representations of high-dimensional data and exploits them to design effective data structures. We also adapt our framework to the problem of estimating frequencies over a data stream, and believe it could be a powerful discovery tool for new problems.

## 1 Introduction

> *Can deep learning models be trained to discover data structures from scratch?*

There are several motivations for this question. The first is scientific. Deep learning models are increasingly performing tasks once considered exclusive to humans, from image recognition and mastering the game of Go to engaging in natural language conversations. Designing data structures and algorithms, along with solving complex math problems, are particularly challenging tasks. They require searching through a vast combinatorial space with a difficult to define structure. It is therefore natural to ask what it would take for deep learning models to solve such problems. There are already promising signs: these models have discovered fast matrix-multiplication algorithms (Fawzi et al., ), solved SAT problems (Selsam et al., ), and learned optimization algorithms for various learning tasks (Garg et al., ; Akyürek et al., ; Fu et al., ; Von Oswald et al., ). In this work, we investigate the problem of data structure discovery, with a focus on nearest neighbor search.

The second motivation is practical. Data structures are ubiquitous objects that enable efficient querying. Traditionally, they have been designed to be worst-case optimal and therefore agnostic to the underlying data and query distributions. However, in many applications there are patterns in these distributions that can be exploited to design more efficient data structures.
This has motivated recent work on learning-augmented data structures which leverages knowledge of the data distribution to modify existing data structures with predictions (Lykouris & Vassilvitskii, ; Ding et al., ; Lin et al., ; Mitzenmacher & Vassilvitskii, ). In much of this work, the goal of the learning algorithm is to learn distributional properties of the data, while the underlying query algorithm/data structure is hand-designed. Though this line of work clearly demonstrates the potential of leveraging distributional information, it still relies on expert knowledge to incorporate learning into such structures. In our work, we ask if it is possible to go one step further and let deep learning models discover entire data structures and query algorithms in an end-to-end manner.

### 1.1 Framework for data structure discovery

Figure: Figure 1: Our model has two components: 1) A data-processing network that transforms raw data into structured data, arranging it for efficient querying and generating additional statistics when given extra space (not shown in the figure). 2) A query-execution network that performs $M$ lookups into the output of the data-processing network in order to retrieve the answer to some query $q$. Each lookup $i$ is managed by a separate query model $Q^{i}$, which takes $q$ and the lookup history $H_{i}$, and outputs a one-hot lookup vector $m_{i}$ indicating the position to query.
Refer to caption: x1.png

Data structure problems can often be decomposed into two steps: 1) data structure construction and 2) query execution. The first step transforms a raw dataset $D$ into a structured database $\hat{D}$, while query-execution performs lookups into $\hat{D}$ to retrieve the answer for some query $q$. The performance of a data structure is typically quantified in terms of two measures: *space complexity*—how much memory is required to store the data structure, and *query complexity*—how many lookups into the data structure are required to answer some query. One can typically tradeoff larger space complexity for smaller query complexity, and vice versa. We focus on these criteria as they are widely studied and have clear practical connections to efficiency.

To learn such data structures, we have a *data-processing network* which learns how to map a raw dataset to a data structure, and a *query network* which learns an algorithm for using the data structure to answer queries (Figure 1). In order to learn efficient data structures and query algorithms we impose constraints on the size of the data structures and on the number of lookups that the query network can make into the data structure. Crucially, we propose end-to-end training of both networks such that the learned data structure and query algorithm are optimized for one another. Moreover, in settings where it is beneficial to learn lower-dimensional representations from high-dimensional data, end-to-end training encourages the representations to capture features of the problem that the data structure can exploit.

On the one hand, learning the data-processing network and query network jointly in an end-to-end fashion seems obvious, especially given the many successes of end-to-end learning over the past decade. On the other hand, it might be hard to imagine such learning getting off the ground. For instance, if the data-processing network produces a random garbled function of the dataset, we cannot hope the query model to do anything meaningful. This is further complicated by the fact that these data structure tasks are more discrete and combinatorial in terms of how the query model accesses the data structure.

### 1.2 Summary of Results

We apply this framework to the problem of nearest neighbor (NN) search in both low and high dimensions. Given the extensive theoretical work on this topic, along with its widespread practical applications, NN search is an ideal starting point for understanding the landscape of end-to-end data structure discovery. Beyond NN search, we explore the problem of frequency estimation in streaming data and discuss other potential applications of this framework. Our findings are:

##### Sorting and searching in 1D (Section 2.2 )

For 1D nearest neighbor search, the data-processing network learns to sort, while the query network simultaneously learns to search over the sorted data. When the data follows a uniform or Zipfian distribution, the query network exploits this structure to outperform binary search. On harder distributions lacking structure, the network adapts by discovering binary search, which is worst-case optimal. Importantly, the model discovers that sorting followed by the appropriate search algorithm is effective for NN search in 1D without explicit supervision for these primitives.

##### K-d trees in 2D (Section 2.3 )

In 2D, when the data is drawn from a uniform distribution, the model discovers a data structure that outperforms k-d trees. On harder distributions, the learned data structure shows surprising resemblance to a k-d tree. This is striking as a k-d tree is a non-trivial data structure, constructed by recursively partitioning the data and finding the median along alternating dimensions.

##### Useful representations in high dimensions (Section 2.4 )

For high-dimensional data, the model learns representations that make NN search efficient. For example, with data from a uniform distribution on a 30-dimensional hypersphere, the model partitions the space by projecting onto a pair of vectors, similar to locality-sensitive hashing. When trained on an extended 3-digit MNIST dataset, the model finds features that capture the relative ordering of the digits, sorts the images using these features, and performs a search on the sorted images—all of which is learned jointly from scratch.

##### Trading off space and query efficiency (Section 2.5 )

An ideal data structure can use extra space to improve query efficiency by storing additional statistics. The learned model demonstrates this behavior, with performance improving monotonically as more space is provided, in both low and high dimensions. Thus, the model learns to effectively trade off space for query efficiency.

##### Beyond nearest neighbor search (Section 3 )

We also explore the classical problem of frequency estimation, where a memory-constrained model observes a stream of items and must approximate the frequency of a query item. The learned structure exploits the underlying data distribution to outperform baselines like CountMin sketch, demonstrating the broader applicability of the framework beyond nearest neighbor search.

## 2 Nearest Neighbor Search

Given a dataset $D=\{x_{1},...,x_{N}\}$ of $N$ points where $x_{i}\in\mathbb{R}^{d}$ and a query $q\in\mathbb{R}^{d}$, the nearest neighbor $y$ of $q$ is defined as $y=\arg\min_{x_{i}\in D}\ dist(x_{i},q)$. We mostly focus on the case where $dist(\cdot)$ corresponds to the Euclidean distance. Our objective is to learn a data structure $\hat{D}$ for $D$ such that given $q$ and a budget of $M$ lookups, we can output a (approximate) nearest neighbor of $q$ by querying at most $M$ elements in $\hat{D}$. When $M\geq N$, $y$ can be trivially recovered via linear search so $\hat{D}=D$ is sufficient. Instead, we are interested in the case when $M\ll N$.(^1^11E.g. in 1D, binary search requires $~{}M=\log(N)$ lookups given a sorted list.)

### 2.1 Setup

##### Data-processing Network

Recall that the role of the data-processing network is to transform a raw dataset into a data structure. The backbone of our data-processing network is an 8-layer transformer model based on the NanoGPT architecture (Karpathy, ). In the case of NN search, we want the data structure to preserve the original inputs and just reorder them appropriately as the answer to the nearest neighbor query should be one of elements in the dataset. The model achieves this by outputting a rank associated with each element in the dataset, which is then used to reorder the elements. More precisely, the transformer takes as input the dataset $D$ and outputs a scalar $o_{i}\in\mathbb{R}$ representing the rank for each point $x_{i}\in D$. These rankings $\{o_{1},...,o_{N}\}$ are then sorted using a differentiable sort function, $sort(\{o_{1},o_{2}\ldots,o_{N}\})$ (Grover et al., ; Cuturi et al., ; Petersen et al., ), which produces a permutation matrix $P$ that encodes the order based on the rankings. By applying $P$ to the input dataset $D$, we obtain $\hat{D}_{P}$, where the input data points are arranged in order of their rankings.
By learning to rank rather than directly outputting the transformed dataset, the transformer avoids the need to reproduce the exact inputs. Note that this division into a ranking model followed by sorting is without loss of generality as the overall model can represent any arbitrary ordering of the inputs.

We also consider scenarios where the data structure can use additional space. To support this use case, the transformer outputs $T$ extra vectors $b_{1},...,b_{T}\in\mathbb{R}^{d}$ which can be retrieved by the query-execution network. We form the data structure $\hat{D}$ by concatenating the permuted inputs and the extra vectors: $\hat{D}=[\hat{D}_{P},b_{1},...,b_{T}]$.

##### Query Execution Network

The role of the query-execution network is to output a (approximate) nearest-neighbor of some query $q$ given a budget of $M$ lookups into the data structure $\hat{D}$. The query-execution network consists of $M$ MLP query models $Q^{1},...,Q^{M}$. The query models do not share weights. Each query model $Q^{i}$ outputs a one-hot vector $m_{i}\in\mathbb{R}^{N+T}$ which represents a lookup position in $\hat{D}$. To execute the lookup, we compute the value $v_{i}$ at the position denoted by $m_{i}$ in $\hat{D}$ as $v_{i}=m_{i}^{\top}\hat{D}$. In addition to the query $q$, each query model $Q^{i}$ also takes as input the query execution history $H_{i}=\{(m_{1},v_{1}),...,(m_{i-1},v_{i-1})\}$ where $H_{1}=\emptyset$. The final answer of the network for the nearest-neighbor query is given by $\hat{y}=m^{\top}_{M}\hat{D}$.

To restrict our model to exactly $M$ lookups, we enforce each lookup vector $m_{i}$ to be a one-hot vector.
Enforcing this constraint during training poses a challenge as it is a non-differentiable operation. Instead, during training, our model outputs soft-lookups where $m_{i}$ is the output of the softmax function and $\sum_{j}m_{ij}=1$. This alone, however, leads to non-sparse queries. To address this, we add noise to the logits before the softmax operation (only during training). This regularizes the query network, encouraging it to produce sparser solutions (see App C.1
for details as to why this occurs). Intuitively, the network learns a function that is robust to noise, and the softmax output becomes robust when the logits are well-separated. Well-separated logits, in turn, lead to sparser solutions. At inference time, we do not add this noise and we ensure the lookup vector $m_{i}$ is a one-hot vector by applying a hardmax function to the network output instead of a softmax.

##### Data Generation and Training

Each training example is a tuple $(D,q,y)$ consisting of a dataset $D$, query $q$, and nearest neighbor $y$ generated as follows: (i) sample dataset $D=\{x_{1},...,x_{N}\}$ from dataset distribution $P_{D}$, (ii) sample query $q$ from query distribution $P_{q}$,
(iii) compute nearest neighbor $y=\arg\min_{x_{i}\in D}dist(x_{i}-q)$. Unless otherwise specified, $dist$ corresponds to the Euclidean distance. The dataset and query distributions $P_{D},P_{q}$ vary across the different settings we consider and are defined later. Given a training example $(D,q,y)$, the data-processing network transforms $D$ into the data structure $\hat{D}$. Subsequently, the query-execution network, conditioned on $q$, queries the data structure to output $\hat{y}$. We use SGD to minimize either the squared loss between $y$ and $\hat{y}$, or the cross-entropy loss between the corresponding vectors encoding their positions.
This is an empirical choice, and in some settings one loss function performs better than the other. All models are trained for at most 4 million gradient steps with early-stopping using a batch size of 1024. After training, we test our model on $10$k inputs $(D,q,y)$ generated in the same way. We describe the exact model architecture and training hyper-parameters in App A.1.

##### Evaluation and Baselines

We evaluate our end-to-end model (referred to as *E2E*) on one-dimensional, two-dimensional, and high-dimensional nearest-neighbor problems. We primarily focus on data structures that do not use extra space, but in Section 2.5, we also explore scenarios with additional space.

We compare against suitable NN data structures in each setting (e.g., sorting followed by binary search in 1D), and also
against several ablations to study the impact of various model components. The *E2E (frozen)* model does not train the data-processing network, relying on rankings generated by the initial weights. The *E2E (no-permute)* model removes the permutation component of the data-processing network so that the transformer has to learn to transform the data points directly. The *E2E (non-adaptive)* model ablation conditions each query model $Q^{i}$ on only the query $q$ and not the query history $H_{i}$. From the $M$ query models, we select the prediction that is closest to the query as the final prediction $\hat{y}$.

### 2.2 One-dimensional data

Figure: Figure 2: (Left) Our model (E2E) trained with 1D data from the uniform distribution over $(-1,1)$ outperforms binary search and several ablations. (Center) Distribution of lookups by the first query model. Unlike binary search, the model does not always start in the middle but rather closer to the query’s likely position in the sorted data. (Right) When trained on data from a “hard” distribution for which the query value does not reveal information about the query’s relative position, the model finds a solution similar to binary search. The figure shows an example of the model performing binary search (’X’ denotes the nearest neighbor location).
Refer to caption: x2.png

##### Uniform Distribution

We consider a setting where the data distribution
$P_{D}$ and query distribution $P_{q}$ correspond to the uniform distribution over $(-1,1)$, $N=100$ and $M=7$. We plot the accuracy(^2^22We include MSE plots as well in App. B.), which refers to zero-one loss in identifying the nearest neighbor, after each lookup in Figure 2 (Left). Recall that $v_{i}$ corresponds to the output of the $i$-th lookup. Let $v_{i}^{*}$ denote the closest element to the query so far: $v^{*}_{i}=\arg\min_{v\in\{v_{1},...,v_{i}\}}||v-q||_{2}^{2}$. At each lookup index we plot the nearest neighbor accuracy corresponding to $v_{i}^{*}$. We do this for all the methods.

A key component in being able to do NN search in 1D is sorting. We observe that the trained model does indeed learn to sort. We verify this by measuring the fraction of inputs that are mapped to the correct position in the sorted order, averaged over multiple datasets. The trained model correctly positions approximately 99.5% of the inputs. This is interesting as the model never received explicit feedback to sort the inputs and figured it out in the end-to-end training. The separate sorting function aids the process, but the model still has to learn to output the correct rankings.

The second key component is the ability to search over the sorted inputs. Here, our model learns a search algorithm that outperforms binary search, which is designed for the worst case. This is because unlike binary search, our model exploits knowledge of the data distribution to start its search closer to the nearest neighbor, similar to interpolation search (Peterson, ). For instance, if the query $q\approx 1$, the model begins its search near the end of the list (Figure 2 (Center)). The minor sorting error ($\sim 0.5\%$) our model makes likely explains its worse performance on the final query.

Figure: Figure 3: For 1D Zipfian query distribution, our model performs slightly better than the the learning-augmented treap algorithm from Hsu et al. () and both methods significantly outperforms binary search.
Refer to caption: x3.png

To understand the relevance of different model components, we compare against various ablations. The *E2E (frozen)* model (untrained transformer) positions only about 9% of inputs correctly, explaining its under-performance. This shows that the transformer must learn to rank the inputs, and that merely using a separate function for sorting the transformer output is insufficient.
The *E2E (non-adaptive)* baseline, lacking query history access, underperforms as it fails to learn adaptive solutions crucial for 1D NN search. The *E2E (no-permute)* ablation does not fully retain inputs and so we do not measure accuracy for this baseline. We verify this by measuring the average minimum distance between each of the transformer’s inputs to its outputs. These ablations highlight the crucial role of both learned orderings and query adaptivity for our model.

##### Zipfian Distribution

Prior work has shown that several real-world query distributions follow a Zipfian trend whereby a few elements are queried far more frequently than others, leading to the development of learning-augmented algorithms aimed at exploiting this (Lin et al., ).
We consider a setting where $P_{D}$ is the discrete uniform distribution over $\{1,...,200\}$ and $P_{q}$ is a Zipfian distribution over $\{1,...,200\}$ skewed towards smaller numbers such that the number $i$ is sampled with probability proportional to $\frac{1}{i^{\alpha}}$. We set $\alpha=1.2$. Again, in this setting $N=100$ and $M=7$.

In Figure 3 we compare our model to both binary search and the learning-augmented treap from Lin et al. (). Our model performs slightly better than the learning-augmented treap and both algorithms significantly outperform binary search with less than $\log(N)$ queries. This setting highlights a crucial difference in spirit between our work and much of the existing work on learning-augmented algorithms. While the Zipfian treap incorporates learning in the algorithm, the authors still had to figure out how an existing data structure could be modified to support learning. On the other hand, by learning end-to-end, our framework altogether removes the need for the human-in-the-loop. This is promising as it could be useful in settings where we lack insight on appropriate data structures. The flip side, however, is that learning-augmented data structures usually come with provable guarantees which are difficult to get when training models in an end-to-end fashion.

##### Hard Distribution

To verify that our model can also learn worst-case optimal algorithms such as binary search, we set $P_{D}$ to be a “hard” distribution, with the property that for any given query there does not exist a strong prior over the position of its nearest neighbor in the sorted data (see App. B.1 for more details). To produce a problem instance, we first sample a dataset from $P_{D}$. We then generate the query by sampling a point (uniformly at random) from this dataset, and adding standard Gaussian noise to it. The hard distribution generates numbers at several scales, and this makes it challenging to train the model with larger $N$. Thus, we use $N=15$ and $M=3$.
In general, we find that training models is easier when there is more structure in the distribution to be exploited.

The model does indeed discover a search algorithm similar to binary search.
In Figure 2 (Right), we show a representative example of the model’s search behavior, resembling binary search (see Figure 16 for more examples).
The error curve in Figure 14 also closely matches that of binary search.

In summary, in all the above settings, starting from scratch, the data-processing network discovers that the optimal way to arrange the data is in sorted order. Simultaneously, the query-execution network learns to efficiently query this sorted data, leveraging the properties of the data distribution.

### 2.3 Two-dimensional data

Figure: Figure 4: Our model’s learned data structure for an instance from the uniform distribution in 2D. While the original order of the stored points showed no structure, the learned data structure arranges points that are close together in the Euclidean plane next to each other.
Refer to caption: x4.png

Figure: Figure 5: The learned data structure resembles a k-d tree in 2D. We show the average pairwise distances (along the first, second, and both dimensions) between points for the learned structure and the k-d tree, with darker colors indicating smaller distances. For the k-d tree, we arrange the points by in-order traversal. It recursively splits the points into two groups based on whether their value is smaller or larger than the median along a given dimension, alternating between dimensions at each level, starting with dimension 1. The learned data structure approximately mirrors this pattern, splitting by dimension 2 followed by dimension 1.
Refer to caption: extracted/5977828/figures/2d_hard_perm.jpg

Beyond one dimension it is less clear how to optimally represent a collection of points as there is no canonical notion of sorting along multiple dimensions. In fact, we observe in these experiments that different data/query distributions lead to altogether different data structures. This reinforces the value in learning both the data structure and query algorithm together end-to-end.

##### Uniform Distribution

We use a setup similar to 1D, sampling both coordinates independently from the uniform distribution on $(-1,1)$. We set $N=100$ and $M=6$, and compare to a k-d tree baseline. A k-d tree is a binary tree for organizing points in k-dimensional space, with each node splitting the space along one of the k axes, cycling through the axes at each tree level. Here, our E2E model achieves an accuracy of $\!~{}75\%$ vs $\!~{}52\%$ for the k-d tree (Fig. 10 in App. B).
The model outperforms the k-d tree as it can exploit distributional information. By studying the permutations, we find that our model learns to put points that are close together in the 2D plane next to each other in the permuted order (see Fig. 4 for an example).

##### Hard Distribution

We also consider the case where we sample both coordinates independently from the hard distribution considered in the 1D setup (see Figure 17 for the corresponding error curve).
We observe that the data structure learned by our model is surprisingly similar to a k-d tree (see Fig 5).
This is striking as a k-d tree is a non-trivial data structure, requiring recursively partitioning the data and finding the median along alternating dimensions at each level of the tree.

### 2.4 High-dimensional data

Figure: Figure 6: (Left) For NN search in higher dimensions (d = 30), the trained models perform comparably to (E2E) or better than (E2E (non-adaptive)) locality-sensitive hashing (LSH) baselines. (Center) When trained with a single query, the model partitions the query space based on projection onto two vectors, similar to LSH. We show the query projection onto the subspace spanned by these vectors and the lookup positions for different queries. (Right) When trained end-to-end to do nearest neighbor search over 3-Digit MNIST Images, our model learns 1D features that capture the relative ordering of the numbers in the images.
Refer to caption: x5.png

High-dimensional NN search poses a challenge for traditional low-dimensional algorithms due to the curse of dimensionality. K-d trees, for instance, can require an exponential number of queries in high dimensions (Kleinberg, ). This has led to the development of approximate NN search methods such as locality sensitive hashing (LSH) which have a milder dependence on $d$ (Andoni et al., ), relying on hash functions that map closer points in the space to the same hash bucket.

We train our model on datasets uniformly sampled from the $d$-dimensional unit hypersphere. The query is sampled to have a fixed inner-product $\rho\in[0,1]$ with a dataset point. When $\rho=1$, the query matches a data point, making regular hashing-based methods sufficient. For $\rho<1$, LSH-based solutions are competitive. We train our model for $\rho=0.8$ and compare it to an LSH baseline when $N=100,M=6$, and $d=30$. The LSH baseline partitions $\mathbb{R}^{30}$ using $K$ random vectors and buckets each point in the dataset according to its signed projection onto each of the $K$ vectors. To retrieve the nearest neighbor of a query point, the baseline maps the query vector to its corresponding bucket and selects the closest vector among $M$ candidates (refer App D for more details).

In Figure 6 (Left), we observe that our model performs competitively with LSH baselines.(^3^33We exclude LSH baselines with larger K as they under-perform.) The non-adaptive model does slightly better as adaptivity is not needed to do well here (e.g., LSH is non-adaptive), and lack of adaptivity likely makes training easier. To better understand the data structure our model learns we consider a smaller setting where $N=8$ and $M=1$. We find that the model learns an LSH like solution, partitioning the space by projecting onto two vectors in $\mathbb{R}^{30}$ (see Figure 6 (Center)).
We provide more details in App C.3.

##### Learning useful representations

High-dimensional data often contains low-dimensional structure, such as data lying on a manifold, which can be leveraged to improve the efficiency of NN search. ML models are particularly well-suited to exploit these structures. Here, we explore whether our end-to-end learning framework can learn representations that capture such structures. This is a challenging task as it involves jointly optimizing the learned representation, data structure, and query algorithm.

We consider the following task: given a dataset of distinct 3-digit handwritten number images, and a query image, find its nearest neighbor in the dataset, which corresponds to the image encoding the closest number to the query image (i.e., *nearest* is defined over the label space).

We generate images of 3-digit numbers by concatenating digits from MNIST (see Figure 13 for image samples). To construct a nearest-neighbor dataset $D$, we sample $N=50$ labels (each label corresponds to a number) uniformly from 0 to 199. For each label, we then sample one of its associated training images from 3-digit MNIST. Additionally, we sample a query label (uniformly over $\{0,..,199\}$) and its corresponding training image and find its nearest neighbor in $D$, which corresponds to the image with the same label. We emphasize that the model has no label supervision but rather only has access to the query’s nearest neighbor. After training, we evaluate the model using the same data generation process but with images sampled from the 3-digit MNIST test set.

As both the data-processing and query-execution networks should operate over the same low-dimensional representation we train a CNN feature model $F_{\phi}$ as well. Our setup remains the same as before except now the data-processing network and query-execution network operate on $\{F_{\phi}(x_{1}),...,F_{\phi}(x_{N})\}$ and $F_{\phi}(q)$, respectively. As the underlying distance metric does not correspond to the Euclidean distance, we minimize the cross-entropy loss instead of the MSE loss. Note that the cross-entropy loss only requires supervision about the nearest neighbor of the query, and does not require the exact metric structure, so it can be used even where the exact metric structure is unknown.

Figure: Figure 7: Trained on 3-digit MNIST images, our data-processing model learns to sort the images without explicit supervision for sorting. While we train our model with datasets of size $N=50$, we show a smaller instance with 5 images for better visualization.
Refer to caption: extracted/5977828/Main/figures/sorted_imgs.png

Ideally, the feature model $F$ should learn 1d features encoding the relative ordering of the numbers, the data model sorts them, and the query model should do some form of interpolation search where it can use the fact that the data distribution is uniform to do better than binary search. This is almost exactly what all models learn to do, from scratch, in an end-to-end fashion, without any explicit supervision about which image encodes which number. In Figure 6 (Right) we plot the learned features of the model. We find that the data model learns to sort the features (Figure 7) with  98% accuracy and the query model finds the nearest neighbor with almost 100% accuracy (Figure 12).

### 2.5 Leveraging Extra Space

Figure: Figure 8: For NN search in 1D the model learns to use extra space and outperforms a bucketing baseline.
Refer to caption: x7.png

The previous experiments demonstrate our model’s ability to learn useful orderings for efficient querying. However, data structures can also store additional pre-computed information to speed up querying. For instance, with infinite extra space, a data structure could store the nearest neighbor for every possible query, enabling $O(1)$ search. Here, we evaluate if our model can effectively use extra space.

We run an experiment where the data and query distribution are uniform over $(-1,1)$ with $N=50,M=2$. We allow the data-processing network to output $T\in\{0,2^{1},2^{2},2^{3},2^{4},2^{5},2^{6},2^{7}\}$ numbers $b_{1},...,b_{T}\in\mathbb{R}$ in addition to the $N$ rankings.
We plot the NN accuracy as a function of $T$ in Figure 6 (Right) compared to a simple bucketing baseline. This baseline partitions $[-1,1]$ into $T$ evenly-sized buckets and in each bucket stores $\operatorname*{arg\,min}_{x_{j}\in D}||x_{j}-l_{i}||$ where $l_{i}$ is the midpoint of the segment corresponding to bucket $i$. The
baseline maps a query to its corresponding bucket and predicts the input stored in that bucket as the nearest-neighbor. Our model’s accuracy monotonically increases with extra space demonstrating that the data-processing network learns to pre-compute useful statistics that enable more efficient querying. We provide some insights into the learned solution in App C.4 and show that our model can be trained to use extra space in the high-dimensional case as well (App C.5).

## 3 Beyond Nearest Neighbor Search

Many other data structure problems beyond nearest neighbor search can be modeled by our framework. Here, we illustrate this broader applicability by applying the framework to the classical problem of *frequency estimation*: a memory-constrained model observes a stream of elements, and is subsequently asked to approximate the number of times a query element has appeared (Cormode & Muthukrishnan, ; Cormode & Hadjieleftheriou, ).
In Section 3.2 we describe several other data structure problems that the framework can be applied to.

### 3.1 Frequency Estimation

Given a sequence of $T$ elements $e^{(1)},...,e^{(T)}\!$ drawn from some universe, the task is to estimate the frequency of a query element $e_{q}$ up until time-step $T$. Specifically, we aim to minimize the mean absolute error(^4^44We use absolute error as this is the metric commonly used in prior work (Cormode & Muthukrishnan, ; Cormode & Hadjieleftheriou, ) but our setup works for squared error as well.) between the true count and the estimated count. As in the nearest neighbor setup, the two constraints of interest are the size of the data structure and the number of lookups for query execution. Consequently, our framework can be easily adapted to model this problem. We also choose this problem to highlight the versatility of our framework as it can be applied to streaming settings.

##### Data processing Network

We model the data structure as a $k$ dimensional vector $\hat{D}$ and use an MLP as the data-processing network which is responsible for writing to $\hat{D}$. When a new element arrives in our stream, we allow the model to update M values in the data structure. Specifically, when an element arrives at time-step $t$, the data-processing network outputs $M$ $k$-dimensional one-hot update position vectors $u_{1},...,u_{M}$ and M corresponding scalar update values $v_{1},...,v_{M}$. We then apply the update, obtaining $\hat{D}_{t+1}=\hat{D_{t}}+\sum_{i=1}^{M}u_{i}*v_{i}$. Unlike in the NN setting where we did not constrain the construction complexity of the data structure, here we have limited each update to the data structure to a budget of $M$ lookups. We do so as in the streaming settings updates typically occur often, so it is less reasonable to consider them as a one-time construction overhead cost.

##### Query processing Network

Query processing is handled in a similar fashion to NN search — we have $M$ query MLP models that output lookup positions. Finally, we also train a MLP predictor network $\psi(v_{1},...,v_{M})$
that takes in the $M$ values retrieved from the lookups and outputs the final prediction.

#### Experiments

Figure: Figure 9: When estimating frequencies of elements drawn from a randomly ordered Zipfian distribution, our model outperforms the CountMinSketch baseline given 1, 2, and 4 queries.
Refer to caption: x8.png

##### Zipfian Distribution

We evaluate our model in a setting where both the stream and query distributions follow a Zipfian distribution. This simulates a common feature of frequency-estimation datasets where a few “heavy hitter” elements are updated or queried more frequently than others (Hsu et al., ). For any fixed training instance, the rank order of the elements in the domain is consistent across both the stream and query distributions, but it is randomized across different training instances. As a result, the model cannot rely on knowing which specific elements are more frequent than others; only the overall Zipfian skew is consistent across training instances.

We use a data structure of size $k=32$ and train our model with $M\in\{1,2,4\}$ queries. Both the data and query distributions are Zipfian over $\{1,...,1000\}$ with a fixed skew of $\alpha=1.2$. We evaluate the mean absolute error over streams of length 100 and compare with the CountMinSketch algorithm, a hashing-based method for frequency estimation (Cormode & Muthukrishnan, ) (See App. E for an overview). Our model’s performance improves with more queries and outperforms CountMinSketch (Figure 9). In this case, CountMinSketch degrades with more queries as for a fixed size memory ($k=32$), it is more effective for this distribution to apply a single hash function over the whole memory than to split the memory into $k$ partitions of size $k/M$ and use separate hash functions. We look at the learned algorithm in more detail and find that our model learns an algorithm similar to CountMinSketch, but with an important difference: it uses an update delta of less than 1 when a new item arrives, instead of the delta of 1 used by CountMinSketch. We find that this can be particularly useful when the size of the data structure is small and collisions are frequent. We hypothesize that the better performance of the learned solution is at least partially due to the smaller delta. In Figure 22, we show that we are able to recover the performance of our learned data structure with 1 and 2 queries when we use a smaller delta.

##### Learning Heavy Hitter Features

In the previous experiment, the Zipfian distribution shape was fixed across training instances but the rank ordering of elements was random. In some settings, however, it may be possible to predict which elements are more likely to occur in the stream. While the exact elements may vary between streams, frequently occurring items could share features across streams. For instance, Hsu et al. () show that in frequency estimation for network flows, certain types of IP addresses receive much more traffic than others. We simulate such a setting by fixing the rank ordering of the Zipfian distribution. However, instead of using a universe of integer elements $\{1,...,K\}$, we instead use their corresponding 3-digit MNIST images with $K=100$ (constructed as in the MNIST NN experiment). Given a stream of integers, we map them to their corresponding MNIST labels and then for each label we sample a corresponding image from the training set. During evaluation, we use images samples from the test set. As the distribution is skewed and the ranking is fixed, images with smaller numbers are sampled much more frequently than those with larger numbers. As in the MNIST NN experiment, we also use a feature-learning CNN model to process the images before passing them to the data-processing and query-execution networks.

We compare our model to CountMinSketch with 1-query that is given the underlying labels instead of the images. Our model has a significantly lower error than the baseline (0.15 vs 2.81 averaged over a stream of size 100 (see Fig. 23)) as the latter is distribution-independent.
By training from the data-distribution end-to-end, our framework is able to simultaneously learn features of heavy hitters (in this case, clustering images with the same label) and use this information to design an efficient frequency estimation data structure. We investigate the learned structure and find that the model has reserved separate memory positions for heavy hitters, thereby preventing collisions (Fig. 24).

### 3.2 Other potential applications

Here, we outline several other potential applications of our framework to facilitate future work.

Graph data structures: Many graph-related problems require an efficient representation to support connectivity or distance queries between vertices. For distance queries, one approach is to use quadratic space to store the distances between all vertex pairs, allowing $O(1)$ query time. Alternatively, one could use no extra space and simply store the graph (which may require significantly less than quadratic space) and run a shortest-path algorithm at query time. The challenge is to find a middle ground: using sub-quadratic space while still answering distance queries faster than a full shortest-path computation (Thorup & Zwick, ).

Sparse matrices:
Another common problem that can be framed as a data structure problem is that of compressing sparse matrices. Given an $M\times N$ matrix, on one hand, one could store the full matrix and access elements in $O(1)$ time. However, depending on the number and distribution of 0s in the matrix, different data structures could be designed that use less than $O(MN)$ space. There is an inherent trade-off between how compressed the representation is and the time required to access elements of the matrix to solve various linear algebraic tasks involving the matrix such as matrix-vector multiplication (Buluç et al., ; Chakraborty et al., ).

Learning statistical models:
Our framework can also handle problems such as learning statistical models like decision trees, where the input to the data-processing network is a training dataset, and the output is a model such as a decision tree. The query algorithm would then access a subset of the model at inference time, such as by doing a traversal on the nodes of the decision tree. This could be used to explore questions around optimal algorithms and heuristics for learning decision tress, which are not properly understood (Blanc et al., ; ).

## 4 Related Work

##### Learning-Augmented Algorithms

Recent work has shown that traditional data structures and algorithms can be made more efficient by learning properties of the underlying data distribution. Kraska et al. () introduced the concept of learned index structures which use ML models to replace traditional index structures in databases, resulting in significant performance improvements for certain query workloads. By learning the cumulative distribution function of the data distribution the model has a stronger prior over where to start the search for a record, which can lead to provable improvements to the query time over non-learned structures (Zeighami & Shahabi, ). Other works augment the data structure with predictions instead of the query algorithm. For example, Lin et al. () use learned frequency estimation oracles to estimate the priority in which elements should be stored in a treap. Perhaps more relevant to the theme of our work is Dong et al. (), which trains neural networks to learn a partitioning of the space for efficient nearest neighbor search using locality sensitive hashing, and the body of work on learned hash functions (Wang et al., ; Sabek et al., ).
While all these works focus on augmenting data structure design with learning, we explore whether data structures can be discovered entirely end-to-end using deep learning.

##### Neural Algorithmic Learners

There is a significant body of work on encoding algorithms into deep networks.
Graves et al. () introduced the Neural Turing Machine (NTM), which uses external memory to learn tasks like sorting and copying. Veličković et al. () used graph neural networks (GNNs) to encode classical algorithms such as breadth-first search. These works train deep networks with a great degree of supervision with the aim of encoding known algorithms. For instance, Graves et al. () use the ground truth sorted list as
supervision to train the model to sort.
There has also been work on learning algorithms in an end-to-end fashion. Fawzi et al. () train a model using reinforcement learning to discover matrix multiplication algorithms, while Selsam et al. () train neural networks to solve SAT problems. Garg et al. () show that transformers can be trained to encode learning algorithms for function classes such as linear functions and decision trees. Our work adds to this line of research on end-to-end learning, focusing on discovering data structures.

## 5 Conclusion

We began with the question of whether deep learning models can be trained to discover data structures from scratch. This work provides initial evidence that it is possible. For both nearest neighbor search and frequency estimation, the models—trained end-to-end—discover distribution-dependent data structures that outperform worst-case baselines. We hope this research inspires further exploration into data structure and algorithm discovery.

One limitation that future research could address is scale. Due to computational constraints, most of our experiments are conducted with datasets of size $N=100$, although in App. F we scale some to $N=500$. While this scale can be sufficient for gaining insights into data structure design, practical end-to-end use would require further scaling. We believe both larger models and better inductive biases could enable scaling up further (see App. F for details).

## 6 Acknowledgements

OS and LC acknowledge the support of the CIFAR AI Chair program. This research was also enabled in part by compute resources provided by Mila – Quebec AI Institute (mila.quebec). GV is supported by a Simons Foundation Investigator Award, NSF award AF-2341890 and UT Austin’s Foundations of ML NSF AI Institute. VS was supported by NSF CAREER Award CCF-2239265 and
an Amazon Research Award.

## Appendix A Training Details

### A.1 Nearest Neighbors

The transformer in the data-processing network is based on the NanoGPT architecture (Karpathy, ) and has 8 layers with 8 heads each and an embedding size of 64. Each query model $Q^{i}$ is a 3-layer MLP with a hidden dimension of size 1024. Each hidden layer consists of a linear mapping followed by LayerNorm (Ba et al., ) and the ReLU activation function (Nair & Hinton, ). In all experiments we use a batch size of 1024, 1e-3 weight decay and the Adam optimizer (Kingma & Ba, ) with default PyTorch (Paszke et al., ) settings. We do a grid search over $\{0.0001,0.00001,0.00005\}$ to find the best learning rate for both models. All models are trained for at most 4 million gradient steps with early-stopping. We apply the Gumbel Softmax (Jang et al., ) with a temperature of $2$ to the lookup vectors to encourage sparsity. All experiments are run on a single NVIDIA RTX8000 GPU.

### A.2 Frequency Estimation

We follow the same setup as the nearest neighbors training except for frequency estimation, the data-processing network is a 3-layer MLP with a hidden dimension of size 1024. We do a grid search over $\{0.0001,0.00005,0.00001\}$ to find the best learning rate for both models. Models are trained for 200k gradient steps with early stopping. All experiments are run on a single NVIDIA RTX8000 GPU.

## Appendix B Additional Nearest Neighbor Performance Plots

Figure: Figure 10: 2D Uniform Accuracy.
Refer to caption: x9.png

Figure: Figure 11: Mean square error plots for (Left) 1D Uniform distribution, (Center) 2D Uniform distribution, (Right) 30D Uniform distribution over unit hyper-sphere.
Refer to caption: x10.png

Figure: Figure 12: 3-Digit MNIST Nearest Neighbors Accuracy. Even though binary search (over the underlying digits) is an unfair comparison, we include it as a reference to compare our model’s performance with.
Refer to caption: x11.png

Figure: Figure 13: Samples from 3-Digit MNIST
Refer to caption: extracted/5977828/Main/figures/mnist_samples.png

### B.1 Hard Distribution

To generate data from the hard distribution, we first sample the element at the 50th percentile from the uniform distribution over a large range. We then sample the 25th and 75th percentile elements from a smaller range and so on. The intuition behind this distribution is to reduce concentration such that $p(NN|q)$ is roughly uniform where $NN$ denotes the index of the nearest-neighbor of $q$ in the sorted list.

Precisely, to sample $N$ points from the hard distribution we generate a random balanced binary tree of size $N$. All vertices are random variables of the form $Uniform(0,a^{\log n-k})$ where $a$ is some constant and $k$ is the level in the tree that the vertice belongs to. If the $i-th$ node in the tree is the left-child of its parent, we generate the point $x_{i}$ as $x_{i}=x_{p(i)}-d_{i}$ where $p(i)$ denotes the parent of the $i-th$ node and $d_{i}$ is a sample from node $i$ of the random binary tree. Similarly, if node $i$ is the right child of its parent, $x_{i}=x_{p(i)}+d_{i}$. For the root element $x_{0}=d_{0}$. In our experiments we set $a=7$. The larger the value of $a$, the greater the degree of anti-concentration. We found it challenging to train models with $N>16$ as the range of values that $x_{i}$ can take increases with $N$. Thus for larger $N$, the model needs to deal with numbers at several scales, making learning challenging.

Figure: Figure 14: Our model’s performance is closely aligned with binary search on the hard distribution in 1D. By design, this distribution does not have a useful prior our model can exploit and so it learns a binary search like solution.
Refer to caption: x12.png

Figure: Figure 15: The positional distribution per lookup in the 1D Hard experiment. Our model closely aligns with binary search, first looking at the middle element, then (approximately) either the 25th or 75th percentile elements, and so on.
Refer to caption: extracted/5977828/figures/1d_hard_lookup_dist.jpg

Figure: Figure 16: Binary Search vs. our model on the hard distribution in 1D.
Refer to caption: extracted/5977828/figures/bs_ex1.jpg

Figure: Figure 17: On the 2D hard distribution our model roughly tracks the performance of a k-d tree.
Refer to caption: x13.png

## Appendix C Additional Experiment Findings

### C.1 Noise Injection for Lookup Sparsity

We find that adding noise prior to applying the soft-max on the lookup vector $m_{i}$ leads to sparser queries. We hypothesize that this is because the noise injection forces the model to learn a noise-robust solution which corresponds to a sparse solution. Consider a simplified setup in 1D where the query model is not conditioned on $q$ and is only allowed one lookup ($M=1$) and $D$ is a sorted list of three elements: $D=[x_{1},x_{2},x_{3}]$. For a given query $q$ and its nearest neighbor $y$, the query-execution network is trying to find the optimal vector $\hat{m}\in\mathbb{R}^{3}$ that minimizes $||y-m^{T}D||_{2}^{2}$ where $m=softmax(\hat{m}+\epsilon),\epsilon\sim$ Gumbel distribution Jang et al. (). Given that $M=1$, the model cannot always make enough queries to identify $y$ and so in the absence of noise the model may try to predict the ’middle’ element by setting $\hat{m}_{1}=\hat{m}_{2}=\hat{m}_{3}$. However, when noise is added to the logits $\hat{m}$ this solution is destabilized. Instead, in the presence of noise, the model can robustly select the middle element by making $\hat{m}_{2}$ much greater than $\hat{m}_{1},\hat{m}_{3}$. We test this intuition by running this experiment for large values of $N$ and find that with noise the average gradient is much larger for $\hat{m}_{N/2}$.

### C.2 2D Uniform Distribution

Figure: Figure 18: k-d search vs. our model on the uniform distribution in 2D. Unlike the k-d tree, our model has a stronger prior over where to begin its search.
Refer to caption: extracted/5977828/figures/2d_uni_ex1.jpg

### C.3 N=8, M=1 30D Experiment

To determine if our model has learned an LSH-like solution, we try to reverse engineer the query model in a simple setting where $N=8$ and $M=1$. The query-execution model is only allowed one lookup. We fit 8 one-vs-rest logistic regression classifiers using queries sampled from the query distribution and the output of the query model (lookup position) as features and labels, respectively. We then do PCA on the set of 8 classifier coefficients. We find that the top 2 principal components explain all of the variance which suggests that the query model’s mapping can be explained by the projection onto these two components. In Figure 19 we plot the projection of queries onto these components and color them based on the position they were assigned by the query model. We do the same for inputs $x_{i}\in D$ and color them by the position they were permuted to. The plot on the right suggests that the data-processing network permutes the input vectors based on their projection onto these two components. This assignment is noisy because there may be multiple inputs in a dataset that map to the same bucket and because the model can only store a permutation, some buckets experience overflow. Similarly, the query model does a lookup in the position that corresponds to the query vector’s bucket. This behaviour suggests the model has learned a locality-sensitive hashing type solution!

Figure: Figure 19: (Left) Projection of queries onto top two PCA components of the decision boundaries of the query model, colored by the lookup position the query is mapped to. (Right) Projection of inputs onto the same PCA components colored by the position the data-processing model places them in. Both the data-processing and query models map similar regions to the same positions, suggesting an LSH-like bucketing solution has been learned.
Refer to caption: extracted/5977828/figures/high_d_interp_old.png

### C.4 1D Extra Space

Figure: Figure 20: (Top) Decision boundary of the first query model. (Bottom) The regression coefficients of the values stored in extra positions as a linear function of the (sorted) inputs.
Refer to caption: extracted/5977828/figures/1d_extra_space_lookup.jpg

#### C.4.1 Bucket Baseline

We create a simple bucket baseline that partitions $[-1,1]$ into $T$ evenly sized buckets. In each bucket $b_{i}$ we store $argmin_{x_{j}\in D}||x_{j}-l_{i}||$ where $l_{i}$ is the midpoint of the segment partitioned in $b_{i}$. This baseline maps a query to its corresponding bucket and predicts the input stored in that bucket as the nearest-neighbor. As $T\to\infty$ this becomes an optimal hashing-like solution.

#### C.4.2 Understanding Extra Space Usage

By analyzing the lookup patterns of the first query model, we can better understand how the model uses extra space. In Figure 20 we plot the decision boundary of the first query model. The plot demonstrates that the model chunks the query space $([-1,1])$ into different buckets. To get a sense of what the model stores in the extra space, we fit a linear function on the sorted inputs and regress the values stored in each of the extra space tokens $b_{i}$ and plot the coefficients for several of the extra spaces in Figure 20. For a given subset of the query range, the value stored at its corresponding extra space is approximately a weighted sum of the values stored at the indices that correspond to the percentile of that query range subset. This is useful information as it tells the model for a given query percentile how ’shifted’ the values in the current dataset stored in the corresponding indices are from model’s prior.

### C.5 30D Extra Space

Figure: Figure 21: Our unconstrained model (E2E) and a more interpretable version (E2E (Coefficients)) both learn to effectively leverage an increasing amount of extra space in 30D, with the unconstrained model outperforming an LSH baseline.
Refer to caption: x15.png

In high-dimensions it is less clear what solutions there are to effectively leverage extra space, and in fact understanding optimal tradeoffs in this case is open theoretically (Arya et al., ).

We follow a similar setup to the 1D extra space experiments but use the data and query distributions from section 2.4. We experiment with two versions of extra space (unrestricted) and (coefficients). For the unrestricted version the data model can store whatever 30 dimensional vector it chooses in each of the extra spaces. For the coefficient model, instead of outputting a 30 dimensional vector, for each extra space, the model outputs a separate N dimensional vector of coefficients. We then take a linear combination of the (permuted) input dataset using these coefficients and store the resulting vector in the corresponding extra positions. While the unrestricted version is more expressive the coefficient version is more interpretable. We include both versions to demonstrate the versatility of our framework. If one is only interested in identifying a strong lower-bound of how well one can use a fixed budget of extra space they may use the unrestricted model. However, if they are more concerned with investigating specific classes of solutions or would like a greater degree of interpretability they can easily augment the model with additional inductive biases such as linear coefficients.

We plot the performance of both models along with an LSH baseline in Figure 21. While both models perform competitively with an LSH baseline and can effectively leverage an increasing amount of extra space, the unrestricted model outperforms the coefficient model at a certain point.

### C.6 Frequency Estimation

Figure: Figure 22: We apply the insight that we learned from our E2E model to improve CountMinSketch on the Zipfian distribution. By changing the CountMinSketch update delta of 1 to our model’s learned delta ($\Delta=0.87$ for $M=1$ and $\Delta=0.93$ for $M=2$), we can improve the performance of CountMinSketch on $M\in\{1,2\}$ queries.
Refer to caption: x16.png

Figure: Figure 23: On the MNIST heavy-hitters frequency estimation experiment, our model significantly outperforms CountMinSketch. This is because our model can learn features predictive of heavy hitters, as opposed to the distribution-agnostic CountMinSketch.
Refer to caption: x17.png

Figure: Figure 24: We show the decision boundary learned by the query/data-processing network in the MNIST heavy hitters experiment. As images with smaller numbers occur more frequently in the stream, the memory-constrained model learns to reserve separate memory positions for these items in order to prevent collisions among them.
Refer to caption: extracted/5977828/Main/figures/cs_mnist_mapping.png

## Appendix D LSH Baseline

Our LSH baseline samples $K$ random vectors $\mathbf{r_{1}},...,\mathbf{r_{K}}$ from the standard normal distribution in $\mathbb{R}^{d}$. For a given vector $\mathbf{v}\in\mathbb{R}^{d}$, its hash code is computed as $hash(\mathbf{v})=[sign(\mathbf{v^{T}r_{1}}),...,sign(\mathbf{v^{T}r_{K}})]$. In total, there are $2^{K}$ possible hash codes. To create a hash table, we assign each hash code a bucket of size $N/2^{K}$. For a given dataset $D=\{x_{1},...,x_{N}\}$, we place each input in its corresponding bucket (determined by its hash code $hash(x_{i})$. If the bucket is full, we place $x_{i}$ in a vacant bucket chosen at random. Given a query $q$ and a budget of $M$ lookups, the baseline retrieves the first $M$ vectors in the bucket corresponding to $hash(q)$. If there are less than $M$ vectors in the bucket, we choose the remaining vectors at random from other buckets. We design this setup like so to closely align with the constraints of our model (i.e. only learning a permutation).

## Appendix E CountMinSketch

CountMinSketch (Cormode & Muthukrishnan, ) is a probabilistic data structure used for estimating the frequency of items in a data stream with sublinear space. It uses a two-dimensional array of counters and multiple independent hash functions to map each item to several buckets. When a new item $x$ arrives, the algorithm computes $d$ hash functions $h_{1}(x),h_{2}(x),\dots,h_{d}(x)$, each of which maps the item to one of $w$ buckets in different rows of the array. The counters in the corresponding buckets are incremented by $1$. To estimate the frequency of an item $x$, the minimum value across all counters $C[1,h_{1}(x)],C[2,h_{2}(x)],\dots,C[d,h_{d}(x)]$ is returned. The sketch guarantees that the estimated frequency $\hat{f}(x)$ of an item $x$ is at least its true frequency $f(x)$, and at most $f(x)+\epsilon N$, where $N$ is the total number of items processed, $\epsilon=\frac{1}{w}$, and $w$ is the width of the sketch. The probability that the estimate exceeds this bound is at most $\delta=\frac{1}{d}$, where $d$ is the depth of the sketch (i.e., the number of hash functions). These guarantees hold even in the presence of hash collisions, providing strong worst-case accuracy with $\mathcal{O}(w\cdot d)$ space.

## Appendix F Limitations and Future Work

Figure: Figure 25: We scale both the 1D (Left) and 30D (Center) experiments to datasets of size $N=500$. (Right) We compare our E2E model with a version where the query-execution network is only composed of one query-model (E2E (shared)) that is used in a loop for $M=7$ queries during training on the 1D Uniform distribution, thereby conserving parameters by reusing weights. This could be a promising direction for problem settings where there is a recursive structure to the query algorithm.
Refer to caption: x18.png

One limitation of our work is the scale at which we learn data structures. Most of our nearest neighbor search experiments are done with input dataset sizes around $N=100$, however, we are also able to scale up to $N=500$ (Figure 25 (Left/Center)), though with less than $\log(N)$ queries. While we demonstrate that useful data structures can still be learned at this scale, it is possible that other classes of structures only emerge for larger datasets. We also believe that many of the insights that can be derived from our models’ learned solutions would scale to larger $N$. For instance, sorting in 1D and locality-sensitive hashing in higher dimensions. We limit ourselves to datasets of these sizes due to computational constraints, and because our primary goal was to understand whether end-to-end data structure design is feasible at any reasonable scale. However, we believe our framework could scale to datasets with thousands of points by increasing the parameter counts of the data-processing and query-execution models. Moreover, as transformers become increasingly efficient at handling larger context sizes in language modeling settings, some of these modeling advancements may also be used for scaling models in the context of data structure discovery.

Complementary to our work, it could also be valuable to explore better inductive biases for the query and data-processing networks, and other methods to ensure sparse lookups, enabling smaller models to scale to larger datasets. For instance, using shared weights among query models can be helpful in scaling up the number of queries. As a first step in this direction we show that a single query model can be used in-a-loop for NN search in 1D (Figure 25 (Right)). However, we leave further investigation for future work.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Q1. Is it possible to verify the effectiveness of the method with a large data that approaches the real applications?
soundness: 2
strengths: S1. This paper studies an interesting problem. 

S2. The method may work. 

S3. It is a novel work.
summary: This paper aims at discovering data structures for nearest neighbor search in an end-to-end manner. The method proposed in this paper employ two networks, i.e., data-processing network and query execution network, for the purpose. The proposed method is pretty costly such that only very few data points are included in the experiments, which makes the method difficult to be evaluated for its practical usefulness.
weaknesses: W1. The method proposed is too costly for finding such data structures. 

W2. The experiments are not sufficient to support its practical usefulness. 

W3. The data structures found seems too simple, i.e., sorting and K-d trees.

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 4
questions: I do not have any particular questions for the authors.
soundness: 3
strengths: The proposed framework is novel as far as I know. It is curious that the models learn to replicate binary search and a k-d tree structure. The article is well-written and easy to follow.
summary: The article proposes an end-to-end learning framework for data structures, and implement it in the cases of sorting, nearest neighbor search, and frequency estimation. The method is based on joint learning of a data processing network (8 layer transformer model in the implementation of the article) and query processing networks (MLPs in the implementation of the article). The data processing network is used to learn an index structure by ranking and then sorting the original data set. The query processing networks are trained to predict the correct output (e.g., the nearest neighbor of the query point) given the output of the data processing network and the query execution history (i.e., the outputs of the earlier lookups). 

The proposed framework differs from the earlier work on learning-augmented data structures since it does not require specifying the data structure that is used, but is completely end-to-end. The authors test they approach on very small ($N \approx 100$, $d =1,2,30$) simulated data sets, and show that it can match the performance of or outperform simple worst-case algorithms and data structures, such as binary search, k-d trees, and locality-sensitive hashing (LSH).
weaknesses: The article is very exploratory. The data sets are simulated and several magnitudes smaller ($N \approx 100$) and have smaller dimensionality ($d=1,2,30$) than the ones used in the practical applications of nearest neighbor search, and the baseline algorithms are very elementary. While the results have some curiosity value, it seems very unlikely that the proposed method could be (a) scaled to the data sets that are used in the real applications of nearest neighbor search (typically, $N>1000 000$ and $d \in [100,1000]$ ) AND (b) match the performance of SOTA approximate nearest neighbor search algorithms, such as ScaNN (Guo et al., 2020) and HNSW (Malkov & Yashunin, 2018). And even if it were be possible, this would require enormous computational resources.

Approximation algorithms, such as approximate nearest neighbor search, are used to _save_ computational resources and speed up ML pipelines that use computationally heavy components. For instance, ANN search has recently been used for approximate attention computation in transformer models (Kitaev et al. 2020; Roy et al., 2021), and to speed up inference in retrieval-augmented generation (RAG) (Borgeaud et al., 2022; Lewis et al., 2020).

Thus, I really fail to see how adding a computationally intensive component, such as a transformer model, to the approximation pipeline could be a useful or valuable contribution, even though it resulted in a slightly decreased query latency compared to the SOTA ANN algorithms (which the article is not even close to demonstrating). In contrast, the SOTA ANN libraries, such as FAISS and ScaNN (Douze et al., 2024) use elementary tools, such as $k$-means clustering, for indexing. This is because they have to scale to billion-scale data sets, and thus the index construction time has to be reasonable.

Borgeaud, Sebastian, et al. "Improving language models by retrieving from trillions of tokens." International conference on machine learning. PMLR, 2022.

Douze, Matthijs, et al. "The faiss library." arXiv preprint arXiv:2401.08281 (2024).

Guo, Ruiqi, et al. "Accelerating large-scale inference with anisotropic vector quantization." International Conference on Machine Learning. PMLR, 2020.

Kitaev, Nikita, Lukasz Kaiser, and Anselm Levskaya. "Reformer: The Efficient Transformer." International Conference on Learning Representations. 2020.

Lewis, Patrick, et al. "Retrieval-augmented generation for knowledge-intensive nlp tasks." Advances in Neural Information Processing Systems 33 (2020): 9459-9474.

Malkov, Yu A., and Dmitry A. Yashunin. "Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs." IEEE transactions on pattern analysis and machine intelligence 42.4 (2018): 824-836.

Roy, Aurko, et al. "Efficient content-based sparse attention with routing transformers." Transactions of the Association for Computational Linguistics 9 (2021): 53-68.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: Do you have code to share to replicate your results?
soundness: 4
strengths: * Impressive visualizations and demonstrations of distribution awareness: the plots showing k-d tree resemblance and effectiveness over binary search on a 1D Zipfian distribution were especially impressive.
* Novel idea with good execution
* Paper is well organized and written, and easy to follow
summary: This paper proposes using deep learning to create data structures. The model learns to rank, i.e. learns a mapping of data element to scalar. The model optionally also can output further auxiliary data to help accelerate the data structure. Querying is performed by M lookups into the data structure, where each lookup is chosen by a separate MLP. The Mth lookup is the query result.

This paper looks at the following data structure problems:
* 1D nearest neighbors
* 2D, 30D, and MNIST (high-dimensional with structure) nearest neighbors
* Frequency estimation, done with a separate model architecture.

The paper shows that the models are effective and adapt to the data distributions presented to it during training.
weaknesses: * Impractical for now: only works for small datasets (100 and 500 elements were tested).
* Still requires a moderate amount of inductive bias: for example, the frequency estimation architecture is very different than the nearest neighbors one, and does not seem obvious to me. It seems that generalizing this approach to new data structure problems isn't entirely trivial and requires at least some trial-and-error, as well as expertise in how classical approaches to the data structure may work.
* In short, the main advantage of this paper's idea is its ability to adapt to data distribution, but the real-world benefit of this is hampered by the fact that for now, this work cannot scale to datasets large enough for distribution-fitting to really show impact (especially when considering the performance impact of MLPs vs, for instance, comparing integers in a BBST).
* No code release, especially disappointing given that the training tasks require only synthetic data and MNIST, and involve small models.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Could the author discuss the justification for why a transformer with $O(n^2)$ complexity is necessary, and why using such a network is still feasible and beneficial for the NN search task? (W1)

2. Could the author demonstrate how well the model performs on real-world data distributions in addition to MNIST by including other real-world datasets that may not follow synthetic patterns? (W2)

3. Could the author clarify the implementation details and/or provide the code for the experiments? (W3)

4. Is there a reason that more recent learning-to-hash methods were not included as baselines for the high-dimensional nearest neighbor search experiment? (W4)

5. The application of the framework to nearest neighbor search and frequency estimation seems somewhat disconnected. Could the authors elaborate on the principles that unify these applications within the framework? Could the authors provide a set of design principles or guidelines for adapting the framework to other tasks? This would make it clearer how researchers can extend this work. (W5)
soundness: 3
strengths: S1. The paper is well-motivated and the problem is interesting. The paper is motivated based on the observation that many data structures are designed to be worst-case optimal and agnostic to the underlying data and query distribution. Neural networks show potential in other tasks e.g. SAT solving. So it is natural to study if neural networks can learn data structures that are optimal for specific data and query distribution in an end-to-end manner.

S2. Experiments show that the proposed framework shows a potential to learn the distribution-specific data structures that are more efficient than the traditional worst-case optimal data structures with synthetic data and the mnist dataset, for the problem of NN search and frequency estimation. This demonstrates the potential of applying neural networks to learn data structures.
summary: This paper presents a framework for end-to-end learning of data structures, enabling deep learning models to autonomously discover efficient data structures optimized for specific data distributions. The framework is applied to two problems: nearest neighbor (NN) search and frequency estimation. Key findings include the model learning sorting and search strategies for NN search that rival or outperform traditional methods like binary search, k-d trees, and locality-sensitive hashing (LSH) in some cases. Beyond NN search, the authors apply the model to frequency estimation in data streams, where it demonstrates competitive performance with traditional data structures like CountMin Sketch.
weaknesses: W1: The paper’s reliance on a transformer model with self-attention for data processing presents efficiency and scalability issues for the NN search task. The self-attention mechanism incurs a computational complexity of $O(n^2)$, limiting its application to datasets with a few hundred points due to memory and processing constraints. Considering the NN search problem has a trivial solution with $O(n)$ complexity, this constraint raises concerns about the framework's practicality in real-world applications where much larger datasets are typical. A more scalable approach, such as a different backbone structure, could significantly extend the applicability of this work by reducing computational demands and enabling a larger $n$. Without such efficiency improvements, the method is impractical.

W2: Most experiments rely on synthetic data, except for the MNIST image dataset, which limits the framework's demonstration of generalizability to learn challenging real-world data distributions. While synthetic data offers control over distributions, testing on real-world datasets with unknown or complex distributions would better illustrate the framework’s ability to discover meaningful data structures. 

W3: The paper’s methodology section lacks sufficient clarity on some technical aspects. 
- (1) The paper uses NanoGPT, a decoder-only transformer that typically requires discrete tokens for input and output. However, the paper doesn’t clearly explain how continuous data inputs are tokenized or how outputs are transformed into the scalar ranking $o_i$ values. More detail on the embedding and decoding process for continuous data is necessary to understand the data flow and model structure. 
- (2) In experiments involving CNNs, essential details about the CNN architecture, such as the number of layers, kernel sizes, or feature extraction strategies, are missing. Including these specifications would make it easier to interpret the model's capabilities and constraints, especially regarding performance in high-dimensional settings.

W4: For high-dimensional NN search, the paper compares its model only against the data-independent LSH baseline. However, recent advancements in unsupervised learning-to-hash methods provide competitive alternatives that also leverage data distribution knowledge, for example [1] and [2]. Including comparisons with these modern hashing approaches could help position the framework’s performance more accurately within the field and reveal any advantages or limitations when competing with other data-aware methods. 

W5: The transition from NN search to frequency estimation is somewhat abrupt and lacks a unified explanation of how these tasks fit within the general framework. The two tasks use different backbones and have different outputs for each component. While both applications involve end-to-end data structure learning, the framework’s flexibility for adapting to different problems is not fully articulated. Clarifying the underlying principles that guide network design for new tasks within the framework would enhance its usability.

References:

[1] Lin, Kevin, et al. "Learning compact binary descriptors with unsupervised deep neural networks." Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.

[2] Jin, Sheng, et al. "Unsupervised semantic deep hashing." Neurocomputing 351 (2019): 19-25.

Meta-review: Thanks for your submission to ICLR.

This paper received four reviews, three of which were on the negative side.  Many of the concerns in this paper revolved around the practicality and utility of the approach, the computational expense, concerns about what was actually discovered, and some issues with limited experimentation.  The positive reviewer felt fairly strongly about the promise of this approach, and was willing to champion the paper.  However, the other three reviewers, even after discussion, still were not enthusiastic enough to push for accepting the paper.

I appreciate the amount of work that the authors went into putting together the rebuttals and responses.  I tend to think that this approach has some promise, and I think with some additional work, this could make a very interesting paper indeed.  But I think that, in its current form, there simply isn't enough enthusiasm from the reviewers to warrant accepting this paper.

Decision: Reject