Discovering Data Structures:
Nearest Neighbor Search and Beyond

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.

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.

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.

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.

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.

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, 2024). 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., 2019; Cuturi et al., 2019; Petersen et al., 2022), 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}]$.

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.

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.

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

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²²2We 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, 1957). 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.

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.

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., 2022b). 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. (2022a). 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.

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.

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

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.

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.

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

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