Discovering Data Structures: Nearest Neighbor Search and Beyond
Authors: Omar Salemohamed
Laurent Charlin
Shivam Garg
Vatsal Sharan
Gregory Valiant
Venue: NeurIPS 2025 poster ()
Keywords: data structures
algorithm discovery
learning-augmented algorithms
Primary area: deep_learning

Abstract: We explore if it is possible to learn data structures end-to-end with neural networks, with a focus on the problem of nearest-neighbor (NN) search. We introduce a framework for data structure discovery, which 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. 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 such as images and exploits them to design effective data structures. Beyond NN search, we believe the framework could be a powerful tool for data structure discovery for other problems and adapt our framework to the problem of estimating frequencies over a data stream. To encourage future work in this direction, we conclude with a discussion on some of the opportunities and remaining challenges of learning data structures end-to-end.

Discovering Data Structures: Nearest Neighbor Search and Beyond
Omar Salemohamed * Université de Montréal, Mila
Laurent Charlin f HEC Montréal, Mila
Shivam Garg' Microsoft Research
Vatsal Sharan University of Southern California
Gregory Valiant Stanford University
Abstract
We explore if it is possible learn data structures end-to-end with neural networks, with focus on the problem of nearest-neighbor (NN) search. To answer this question, we introduce framework for data structure discovery which adapts to the underlying data distribution and provides fine-graiee 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. In several settings, we are able reverse-engineer the learned data structures and query algorithms. For ID nearest neighbor search, the model discovers optimal distribution (in)dependent algorithms such as binary search and variants of inter- polation search. In higher dimensions, the model learns solutions that resemble k-d trees some regimes, while in others, elements of locality- emerge. Additionally, the model learns useful representations of high- -igg--iinnsional data such as images and exploits them to design effective data structures. Beyond NN search, believe the framework could be powerful tool for data structure discovery for other problems, and adapt it to the problem of estimating frequencies over data stream. To encourage future work in this direction, we conclude with discussion on some of the opportunities and remaining challenges of learning data structures end-to-end!
hashing
Introduction
Can neural networks discover data structures from scratch?
There several motivations for this question. The first is scientific. Deep learning models 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 combinatorial space with difficult to define structure. is therefore
require searching through natural to ask what would take for deep learning models to solve such problems. There are already promising signs: these models have discovered fast matrix multiiiimutippication algorithms [I], solved SAT problems 22 and learned optimization algorithms for various learning tasks 3 5)
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
Cooreeponndig Author: malemooeeee gemeemooohamme@@@milcccmmm gmail.com
Authors listed in alphabetical order.
3 Our code is available at: http:://ttttttoooococcaasssassse
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

work on learning-ause augmented data structures which leverages knowledge of the data distribution to modify existing data structures with predictions 6
g 9 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 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. is natural to ask we can go a step further and let deep learning models discover entire data structures and query algorithms in an end-to-end manner.
7
In this work, we investigate the possibility of end-to-ent data structure discovery, with focus on the nearest neighbor search problem. A priori. it is not obvious how to frame data structure discovery as an end-to-end learning problem. What are the underlying principles that unify different data structure problems? How can a data structure be represented and queried using neural networks? How do we ensure some notion of efficiency is enforced? Thus a main contribution of our work is in proposing how to frame the problem. To help chart the landscape of this paradigm and explore its scope and limits, we focus more on understanding the kind of algorithms that end-to-end learning discovers than results any given benchmark.
1.1 Framework for data structure discovery
Data structure problems are often divided into two steps: 1) construction and 2) query execu- tion. The first step transforms a raw dataset D into a structured database D, while the sec- ond performs lookups into D to answer a query g. The performance of data structure is typi- cally quantified in terms of two measures: space complexity- the memory required to store the structure, and query complexity- the number of lookups needed to answer query. One can typically tradeoff larger space complexity for smaller query complexity, and vice versa. We focus on these criteria as they are well-studied and directly impact practical efficiency.
Desaeigoceeinng Network
Pb
Datt-prorsiing
2) Query Exeration
Figure 1: Our model has two components: 1) A data-processing network transforms raw data into structured data, arranging it for efficient query- ing and generating additional statistics when given extra space (not shown in the figure). 2) A query- exccution network performs M lookups into the output of the data-proesssing network retrieve the answer to some query q. Each lookup is man- aged by a separate query model Q' which takes g and the lookup history H,, and outputs a one-hot lookup vector m, indicating the position query.
To learn data structures, we have a data- processing network that learns how to map raw dataset to data structure, and a query network that learns an algorithm to answer queries using the data structure (Fig. []. To ensure efficiency, we impose constraints on the data structure's size and the number of lookups the query net- work makes. Crucially, we propose end-to-end training of both networks such that the learned data structure and query algorithm are optimized for one another. In settings where it is beneficial to learn lower-dimensional representations from high-dimensional data, end-to-end training also encourages the representations to capture features that the data structure can exploit.
Why should learning data-structures end-to-end be possible? On one hand, jointly learning the data- processing and query networks end-to-end seems obvious, given the many successes of end-to-end learning over the past decade. On the other hand, it can be hard to imagine such learning getting off the ground. For instance, if the data- aatt--pocessing network produces a random garbled function of the dataset, the query model cannot do anything meaningful. In fact, do observe cases where end-to-end learning fails. This challenge further compounded by the discrete and combinatorial nature of how the query model accesses the data structure and the minimal supervision provided to the models. Despite these challenges, we find that in several settings end-to-end learning can recover variety of classical algorithms.
1.2 Summary of Results
We focus our investigation on the problem of nearest neighbor (NN) search in both low and high dimensions. NN search is an ideal starting point for understanding the landscape of end-to-end data structure discovery given the extensive theoretical work on this topie, along with its widespread practical applications (e.g., similarity search over language/image embeddings [0)). Moreover, the

---

NN search problem has a rich design space and remains an active area of research [1] [3]. In addition NN search, we also explore the problem of frequency estimation streaming data and discuss other potential applications for end-to-end data structure learning (Section
Our findings are:
2.3 For ID 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 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 ID without explicit supervision for these primitives. Similarly, in 2D, the model learns data structure that outperforms k-d trees for uniform distributions and demonstrates recursive partitioning behavior, resembling k-d trees on harder distributions, by constructing medians along alternating dimensions.
Learning Classical Algorithms (Sections 2.2
Useful representations in high dimensions (Section 2.4) For high-dimensional data, the model learns representations that improve NN search efficiency. For instance, with data from a uniform distribution 30-dimensional hypersphere, it partitions space by projecting onto pair of vectors, similar to locality-sensitive hashing. When trained on an extended 3-digit MNIST dataset, the model identifies features that capture number order, sorts images accordingly, and searches within the sorted set set-all learned jointly from scratch! Moreover, it performs competitively with existing data structures datasets like FashionMNIST and SIFT.
Beyond NN search (Section 3) To demonstrate broader applicability of the end-to-end learning paradigm, we study frequency estimation, classic problem where nemooy-constrained model observes a stream and estimates a query item's frequency. The learned structure leverages the data distribution to outperform baselines like CountMinSketch. Moreover, we use insights from the learned model to improve CountMinSketch on a practical dataset. We also outline several other problems for future work.
In summary, we take a first step toward end-to-end data structure learning, showing that it is not only feasible, but also capable of exploring a rich algorithmic space and recovering meaningful- sometimes even elassical solutions from scratch.
At the same time, our work has important limitations. Our experiments are limited to small-scale settings, and further work needed scale them up. Also, while we focus on enforcing query and space complexity constraitss-arguabe such preprocessing time and the inference cost of neural networks require more investigation. We discuss these limitations and potential solutions in Section 5
the most classical measures of efficiency-other aspects,
2 Nearest Neighbor Search
Given a dataset D {21,.., vn of N points where 2; E Rd and query q E Rd the nearest neighbor y of is defined as = arg minl,eD disl(x., q). We mostly focus on the case where dist(.) corresponds to Euclidean distance. Our objective is to learn a data structure Á for D such that given & and budget of M lookups, we can output a (approximate) nearest neighbor of q by querying at most M elements D. When M 2 N, y can be trivially recovered via linear search so D D is sufficient. Instead, we are interested in the case when M < N
2.1 Setup
Data-processing Network Recall that the role of the data- protea-ronssing network transform raw dataset into data structure. The backbone of our data cata-processing network is an 8-layer transformer model based on the NanoGPT architecture 14] (see App B. for all architecture details). We use quadratic-attention
transformer to keep the data- pata-processing model relatively general, however, we find that in certain settings can be possible to use a cheaper alternative such as transformer with linear attention
[5] See App. D.T for more detailed discussion.
g. in ID, binary search requires M log(NV lookups.
3

---

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 rank associated with each element in the dataset, which then used to reorder the elements. More precisely, the transformer takes as input the dataset D and outputs scalar O, E I representing the rank for each point I E D. These rankings for,..., ON} f are then sorted using a differentiable sort function, sort( (01, 02 ....... oN [66] (7[18). which produces permutation matrix P that encodes the order based on the rankings. By applying P to the input dataset D, we obtain Dp, where the input data points arranged in order of their rankings. By learning to rank rather than directly outputting the transformed dataset, the transformer avoids the need 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. See App. E.2 for more information why we learn permutations.
Query Execution Network The role of the query-execution network is to output a nearest-neighhet of a query Q given a budget of M lookups into the data structure î. This introduces a combinatorial constraint-only a fixed number of discrete memory accesses are allowed-while we also require differentiability for gradient-base
training, creating a tension between discreteness and optimization.
To implement this, the network consists of M MLP query models Q',..,, QM, The query models do not share weights to keep the model relatively general, but we explore shared weights in App D.I Each query model Q' outputs a one-hot vector m, E R which represents lookup position in D. To execute the lookup, we compute the value v at the position denoted by m; in π as ey m. D. In addition to the query q, each query model Q' also takes as input the query execution history H; = {(m11,1)) ... (mi-1, 1)} where Hi = 0. The final answer of the network for the nearest- query is given by ÿ = mas D. To enforce exactly M lookups while maintaining differentiability, we train with coftmax-based soft lookups but add noise to their logits. We find this noise encourages sparser logits during training. At inference, we replace softmax with hardmax to produce exact one-hot vectors. In App. B.3 we provide some intuition for why we think adding noise to the logits induces sparsity.
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 {01 {..,..., ..., NN\ from dataset distribution Pp. (ii) sample query q from query distribution Pg. (iii) compute nearest neighbor y = arg min,, e dist(x; - q). Unless otherwise specified, dist corresponds to the Euclidean distance. The dataset and query distributions Pp. Pg vary across the different settings we consider and are defined later. Given a training example (D, q, y), the data-ppocessing network transforms D into the data structure D. Subsequently, the query-execution network, conditioned on q, queries the data structure to output ÿ. We use SGD minimize either the squared loss between y and ÿ. 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 2 million gradient steps with early-stopping using a batch size of 1024. After training, we test our model on 10k inputs (D, 4. 3) generated in the same way. See App B.1 for more details.
Evaluation and Baselines We evaluate our end-to-end model (referred to as E2E) on I-dimensional, 2-dimensional, and high-dimensional NN search. We primarily focus on data structures that do not use extra space, but we also explore scenarios with additional space in App B.
We compare against suitable NN data structures in each setting (e.g., sorting followed by binary search in ID), and also against several ablations to study the impact of various model components. The E2E (frozen) model does not train the data- proce-soosssing 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 conditions each query model Q' on only the query q and not the query history H;. We select the prediction that is closest to the query as the final prediction ÿ.
2.2 One-dimensional data
Uniform Distribution We consider a setting where the data distribution Pp and query distribution
P, correspond to the uniform distribution over (-1, 1), N = 100 and M = 7. We plot the accuracy, which refers to zero-one loss in identifying the nearest neighbor, after each lookup in Fig. 2 (Left)

---

1.0 Nodel
840
0.3
(frizon)
Lookup Position
8 56 8 8
1000
Position
Accuracy
0035
0.2
0.0
Lookup Index
Query
Lookup Index
Figure 2: (Left)) Our model (E2E) trained with ID data from the uniform distribution over (-1,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 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 '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). See Fig. 13 for more examples.
(we include MSE plots well in App. B.4, Recall that v, corresponds the output of the i-th lookup. Let oa be the closest element to the query so far: 2z = arg m ||0 g11.. For all the methods, we plot the nearest neighbor accuracy corresponding
of for each lookup index
key component of ID NN search is sorting, and we observe that the trained model does indeed learn to sort effectively. We verify this by measuring the fraction of inputs mapped to the correct position in the sorted order, averaging 99.5% accuracy across multiple datasets. Remarkably, the model achieves this without explicit feedback to sort, learning the behavior through end-to-end training. While the separate sorting function aids this process, 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 search algorithm that outperforms binary search, which is designed for the worst case. This is because unlike binary search, model exploits knowledge of the data distribution to start its search closer to the nearest neighbor, similar to interpolation search [9]. For instance, the query R3 1, the model begins its search our model makes likely explains its worse performance on the final query.
the end of the list (Fig. 2 (Center)). The minor sorting error (~ 0.5%)
To understand the relevance of different model components, we compare against various ablations. The E2E (fwzzn) model (untrained transformer) positions only about 9% of inputs correctly, ex- plaining 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 insufficient. The E2E (non-adaptive) bascline, lacking query history access, underperforms solutions crucial for ID NN scarch. The E2E (no-permute) ablation does not fully retain inputs and 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 model.
it fails to learn adaptive
Hard Distribution To verify that our model can also learn worst-case optimal algorithms such as binary search, we set Pn to a "hard" distribution, where for any query, no strong prior exists over the position of its nearest neighbor in the sorted data (see App. fo6 more details). To produce a problem instance, we first sample a dataset from Pp. 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 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 search algorithm similar to binary search. In Fig. 2 (Right), show a representative example of the model's search behavior, resembling binary search (see Fig. [3 for more examples). The error curve in Fig. a closely matches that of binary search.
Along with uniform and hard distributions, we show that end-to-end learning is possible also with Zipfian distribution, which is common in many practical settings (see App. B.7).
5

---

Learned Data Structure
K-d Tree Dimersion 2
Dimension 1
Dimension 2
Both Dimensions
Dimemsion
Both Dimenssoos
2
2
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 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.
Figure 3:
Dataset
LSH
E2E
ITQ
K.Means
NeuralLSH
Hypersphere
30.0
35.0
34.4
35.3
35.1
PashionMNIST
30.2
67.2
42.7
66.2
68.2
SIFT
14.3
46.9
30.1
47.1
46.7
MNIST (3-digit)
76.2
97.7
Table 1: NN search accuracy high-d datasets compared to Decality-sensitvee hashing and learning- to-hash baselines. For a description of the baselines and a discussion on why we only compare with LSH on the MNIST dataset, App. B.8
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 efficiently query this sorted data, leveraging the properties of the data distribution.
2.3 Two-dimensional data
Beyond one dimension 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 setup similar 1D, sampling both coordinates independently from the uniform distribution on 11 We set = 100 and M = 6, and compare to a k-d tree baseline. k-d tree is binary tree for organizing points k-dimensional space, with each node splitting the space along one of the k axes, cycling through the axes at each tree level. Here, the E2E model achieves an accuracy of 75% 52% for the k-d tree (Fig. 7 App. B.4]. The model outperforms the k-d tree as it can exploit distributional information. By studying the permutations, we find that the model learns to put points that are close together in the 2D plane next to each other in the permuted order (see Fig. 99 for an example).
Hard Distribution We also consider the case where both coordinates are sampled independently from the hard distribution used in the ID setup (see Fig. [4 for the error curve). We observe that the data structure learned by the model is surprisingly similar to k-d tree (see Fig 3). This is striking as a k-d tree is 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-D data (lypershev/FsshoonNNNSTAYY'
High-dimensional NN search poses 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 [20). This has led to the development of approximate NN search methods such

---

Raw Datasat
9 7 0
102
PCAZ
Leamedd
I Learned Data Stracture
O 2 4
102
I 9 7
PCA1
Figure 4: (Left) On the high-dimensional hypersphere, when trained with single query, the model partitions the query space based on a 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, where position is encoded by color. (Center) When trained end-to-end to do nearest neighbor search over 3-digit MNIST images, model learns ID features that capture the relative ordering of the numbers. (Right) Trained on 3-digit MNIST images, our data- proccessnn model learns to sort the images without explicit supervision for sorting. While train our model with datasets of size N = 50, we show a smaller instance with 5 images for better visualization.
as locality sensitive hashing (LSH) which have milder dependence dimension d (1) 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 fixed inner-product p E [0. 1] with a dataset point. When p 1, the query matches data point, making regular hashing- mashing-based methods sufficient. For < 1, LSH-- based solutions are competitive. We train our model for p = 0.8 and compare it to LSH [22] and Beaming-to-hass baselines when N 100, .M- = 6, and d 30 (see App. B.8 for bascline details).
We observe that our model slightly outperforms the LSH baseline and performs competitively with the tearning-to-hanl baselines (Table D. To better understand the data structure the model learns we consider smaller setting where N and M
We find that the model learns an LSH- like solution, partitioning the space by projecting onto two vectors in R:30 (see Fig. 4 (Left)). Specifically, both the pata-srocsssing and query models learn the same partitioning of space in tandem. These partitions serve the same role as hash functions in LSH algorithms like SimHash [22). Like LSH algorithms, to find a nearest-neighhor, the query model maps a point to corresponding hash bucket and then compares it to other points in that bucket, ultimately selecting the closest one. We provide more details in App B.9
We include additional high-dimensional (100D) experiments on two standard approximate NN benchmarks, FasHionMNIST and SIFT [10, demonstrate the model's performance on realistic data (see App 1.10 A. more details). Our model performs competitively with learning-tor -o-hhh baselines (Table Lt. illustrating that E2E learning can recover reasonable solutions in a variety of settings, even when compared to carefully hand hand-designed solutions. Note that we do not expect E2E to surpass learning-to tomning-to-haste tosh baselines these settings since query adaptivity does not appear to be helpful for these problems. Next, we consider setting where query adaptivity is beneficial.
Learning useful representations (MNIST) digheddimnsionnl data often contains diw-dimensional structure which can be leveraged to improve the efficiency of NN search. Here, we explore whether end-to-end framework learn representations that capture such struetures, challenging task that requires jointly optimizing the learned representation, data structure, and query algorithm.
We consider the following task: given dataset of distinct 3-digit MNIST images (Fig 20) and a query image, find its nearest neighbor-define as the image encoding the closest number to the query (i.c., nearest in label space). Instead of operating directly on pixel space, the data and query networks use low-dimensional representations learned by a CNN feature model F (see App Pop for details).
Ideally, the feature model F should learn 1d features encoding numerical order, the data model sorts them, and the query model performs a form of interpolation search using the fact that the data distribution is uniform to outperform binary search. This is almost exactly what is learned end-to-ent end, from scratch, without any explicit supervision about which image encodes which number. In Fig. (Center) we plot the learned features. We find that the data model learns to sort the features (Fig- (Right)) with 98% accuracy and the query model finds the nearest neighbor with Rs 98% accuracy (Table ID. Moreover, E2E learning beats LSH basclines query adaptivity is useful here.
7

---

Notably, unlike hand-designn data structures that typically assume access to standard distance metrics (e.g., Euclidean). this approach requires no supervision on the underlying metric structure. The model only needs supervision on which dataset element is the nearest neighbor of a query, making it applicable even when the distance metric is unknown or implicit.
Leveraging Extra Space Beyond learning effective orderings for querying, our framework can also learn to use extra memory to accelerate search. In App. B.12 we show that in both low and high-dimensions the data model can learn to pre-compute and store helpful statistics in extra space that enables the query model to find the nearest neighbor with less queries.
Beyond Nearest Neighbor Search
Next, we illustrate the broader applicability of the end-to-end learning paradigm by applying it to the classical problem frequency estimation in the streaming setting. We then describe several other problems in Section 3.2
that this paradigm can be applied to.
3.1 Frequency Estimation
Queries
Given stream Tf elements e(2), e from universe, the task to estimate the frequency of query element E, up to time-step T, aiming to minimize the mean absolute error between true and estimated counts. As in the NN setup, the key constraints are data structure size and the number of lookups for query execution, making this problem compatible with our framework. We explore frequency estimation second task mainly because it is structurally very different from NN search due the streaming nature, yet E2E learning still works. Due to space constraints, we mainly discuss our findings here and leave implementation details to App C.2 We also discuss how both NN search and frequency estimation fit into our broader framework in Apple
Mean Absolute Error
Model EZE
SountMinStetch
20
40
60
80
100
Stream Index
Best Delta Emr/Detta11 Err
We evaluate our model in setting where both stream and query distributions follow Zipfian distribution, simulating frequency- estimation datasets where a few "heavy hitter" elements are queried more frequently 23 For given training sample, the rank order of elements is consistent across both stream and query distributions but randomized across different training samples. Consequently, the model cannot rely on specific elements being more frequent- only overall Zipfian skew consistent.
1.5
1.4
2043 8192
Memory
Figure 5: (Top) When timating frequencies of ele- ments drawn from randomly ordered Zipfian distribution, our model outperforms the CMS baseline given 1, 2, and queries. (Bottom) By aug- menting CMS with an update delta (A) 1 we can out- perform vanilla CMS on the CAIDA IP dataset. For each memory size, we plot the rela- tive performance the best vs the default, A 1.
We compare our model with CountMinSketch (CMS), a hashing- based algorithm for frequency estimation 244 (See App. C.1 for an overview). Our model's 's performance improves with more queries and outperforms CMS (Figure 5 We find that our model learns an algorithm similar to CMS but uses smaller update delta (6 which we hypothesize improves performance when collisions are trequent.
We use this insight to design modified version CMS that uses custom update delta. In Fig. [ show that this augmented CMS algorithm can outperform vanilla CMS by up to factor of Rv 1.9x on real IP traffic dataset 25 consisting of stream of 30 million IP addresses. These results demonstrate learning data structures end-to-end can provide useful insight into data structure design that can be transferred to realistic settings (see App C.3 for more details and App C.4 for additional experiments on frequency estimation nigh-dimensional settings).
*CMS degrades with more queries as for fixed size memory (k = 32). it is more effective for this distribution apply single hash function the whole memory than to split the memory into k partitions of size k/M and use separate hash functions.
'The update delta is the scalar increment value used to update the count of an item in the CMS hash table. For CMS, this value 1. See C.3 for overview of the CMS algorithm.
8

---

3.2 Other potential applications
Our framework is designed for problems that require efficiently answering queries over some col- lection of data. We believe it can be applied to any problem that shares this structure-namely, one where there is a data-processing step, a query algorithm, and constraints on space and query-time complexity. Here, we outline several such candidate problems that could benefit from this approach. We provide more detailed discussion on the broader applicability of the framework and how it can be adapted these problems in App.
Graph data structures Efficient graph representations are essential for connectivity or distance queries. Storing all distances between vertices in quadratic space enables (1)) queries, but requires substantial memory. An alternative is storing the entire graph and running a shortest-par path algorithm query time. The challenge is finding balance: using sub-quadratic space while still answering queries faster than full shortest-path computations [26].
Sparse matrices Another problem that can be framed as data structure problem is compressing sparse matrices. Given an M x N matrix, on one hand, one can store the full matrix and access elements in ((1) time. However, depending structures that use less than O(MN) space could be designed. 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 [27 28
the number and distribution of Os in the matrix, data
Learning statistical models Our framework can handle problems such as learning statistical models, where the input to the data-processing network is a training dataset, and the output 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 optimal algorithms and heuristics for learning decision trees, which are not properly understood [29] 30).
4 Related Work
Learning-Augmmettrr Algorithms Recent work has shown that traditional data structures and algorithms can be made more efficient by learning properties of the underlying data distribution. 31 introduced the concept of learned index structures which use MI. 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 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 [32]. Other works augment the data structure with predictions instead of the query algorithm. For example, [8] use learned frequency estimation oracles to estimate the priority in which elements should be stored in treap. Perhaps more relevant to the theme of our work is [33] which trains neural networks to learn partitioning of the space for efficient nearest neighbor search using locality sensitive hashing, and the body of work on learned There is also work on learning the parameters of nearest neighbor search algorithms using machine learning
hash functions 34 35
36 as well as works that have explored Barning-augmenter 23 37 While 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. Our approach eliminates the human-in-the-lod limited insights into suitable data structures. However, this comes at the cost of losing the provable methods typically offer.
algorithms for frequency estimation
making it promising for settings with
guarantees that
Neural Algorithmic Learners There is significant body of work on encoding and learning algorithms with neural networks. Graves et al. (38) introduced the Neural Turing Machine which uses external memory to learn tasks such as sorting. The neural algorithmic reasoning line of work [39 [0] 41044122 44) aims to simulate a variety of classical algorithms using neural networks. These works typically train models with intermediate outputs of the algorithm they are trying to learn whereas we rely on minimal supervision. 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. [2] train neural networks to solve SAT problems. Garg et al. [3 show that transformers can be trained to encode learning algorithms for function classes such as linear functions and decision trees. A recent line of work on neural sketching algorithms 43 4 shares several similarities to our frequency estimation experiments. Both these methods and ours can be viewed as memory- augmented neural networks that learn to read and write from a differentiable
9

---

memory. This line of work is more focused on developing scalable sketching algorithms whereas we explore whether or not neural networks can discover such algorithms from scratch with minimal inductive biases. See App C.6 for additional discussion.
5 Discussion
Limitations While our work takes first step in exploring end-to-end data structure discovery. there are several important limitations.
Efficiency: Our investigation focuses on the space and query efficiency of the learned data structure. However, in practice, factors such as the neural net inference cost and the pre-processing time to construct the data structure are also important. For example, using standard Transformers in the data model leads to pre-processing time that scales quadratically with the number of items. In preliminary experiments, we find that substituting the quadratic attention transformer with a more efficient alternative, such as linear attention (App. D.I. is often feasible. Further exploration of these architectural choices an important direction for future work. We also benchmark the computational overhead of our methods in App D.3
Scale: Due computational constraints, most of our experiments are conducted on datasets with N = =100 (with some scaled to N = 500; see App. D.2p. Practical end-to-end deployment would require further scaling, which remains an open challenge. While increasing both the model sizes and training time will help to scale to larger datasets to some extent, it will also be important to explore new ideas as well such as better inductive biases that enable smaller models handle large datasets more effectively. For instance, sharing weights among query models may help scale the number of queries. See App. D.2 for more detailed discussion on scaling and preliminary experiments.
In addition to the above limitations, learned data structures lack the same level of interpretability and provable guarantees as their classical counterparts. However, in practical settings, the benefits of learned distributtonndependden algorithms may outweigh these limitations.
Future Applications While scaling is necessary for practical deployment, small-scale models can still yield valuable insights that generalize to larger settings. For example, an Al-in-the-loop discovery workflow might: (1) train model on target problem to surface promising heuristics (c.g.. space partitioning criteria), (2) analyze its behavior to extract actionable insights, and (3) validate those insights theoretically
empirically. Such workflows are feasible Section [] for instance, we use insights from our model to design an augmented version of CMS that outperforms the original algorithm on practical dataset. Moreover, the ability of E2E learning to recover classical algorithms such as binary search, k-d trees, and LSH shows that these models can effectively explore the data structure design space, highlighting the promise of this approach.
with small-scale models. In
We set out by asking: can neural networks discover data structures? At the outset, many basic questions were unclear. How should data structures and queries be represented using neural networks-especialll
Conclusion
under combinatorial constraints like limited query complexity? Even with reasonable setup, it was not clear if training would take off, as both networks are trained from scratch and depend on each other for useful signal. And even if learning does take off, can it explore the rich space of possible data structures and algorithms?
To our surprise, once we set up the training with space/query constraints and a structured lookup interface the model did learn meaningful algorithms. It recovered classic data structures like sorting, k-d trees, and LSH-like behavior, often adapting them to the data distribution. In more complex cases like the extended MNIST experiment, learned image features, built data structure over them, and learned query network to search it-all from scratch and with minimal supervision.
We view this as the main contribution of the paper: framing data structure discovery learning problem and showing that not only works, but can explore a rich solution space. We hope this research encourages further exploration of end-to-end learning for data structures- ferettrresf-rmmmm tackling efficiency and scaling challenges to applying the framework to new problems and leveraging it as an Al-in-the-lon tool for discovery.
end-to-end
10

---

Acknowledgments
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 Al Institute (mila.quebec). GV is supported by a Simons Foundation Investigator Award, NSF award AF-2341899 and UT Austin's Foundations of ML NSF AI Institute. VS was supported by NSF CAREER Award CCF-2239265, an Amazon Research Award, Google Research Scholar Award and an Okawa Foundation Research Grant. This work was done in part when VS was visiting the Simons Institute for the Theory of Computing.
References
[1] Alhussein Fawzi, Matej Balog. Aja Huang. Thomas Hubert, Bernardino Comera-Paredess Mohammadamin Barckatain, Alexander Novikov, Francisco J R Ruiz, Julian Schrittwieser, Grzegorz Swirszez, et al. Discovering faster matrix multiplication algorithms with reinforcement learning. Nature, 610(7930):47-555 2022.
[2] Daniel Selsam, Matthew Lamm, Benedikt Bünz, Percy Liang, Leonardo de Moura, and David Dill. Learning SAT solver from single-bit supervision. arXiv preprint arXiv: arXiv:/802 /ttv./800.000 1802.33685 2018.
[3] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598 2022.
[4] Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint arXiv:2310.17088. 2023.
[5] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers lear in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174,
PMLR, 2023.
[6] Thodoris Lykouris and Sergei Vassilvitskii. Better caching with machine learned advice. 2018.
[7] Jialin Ding, Umar Farooq Minhas, Jia Yu, Chi Wang, Jaeyoung Do. Yinan Li, Hantian Zhang, Badrish Chandramouli, Johannes Gehrke, Donald Kossmann, al. Alex: updatable adap- tive learned index. In Proceedingy of the 2020 ACM SIGMOD International Conference on Management of Data, pages 969-984, 2020.
[8] Honghao Lin, Tian Luo, and David Woodruff. Learning augmented binary search trees. In International Conference on Machine Learning, pages 13431-13440. PMLR, 2022.
[9] Michael Mitzenmacher and Sergei Vassilvitskii. Algorithms with predictions. Communications of the ACM, 65(7):33-335 2022.
[10] Martin Aumüller, Erik Bernhardsson, and Alexander Faithfull. Ann- han-benchmarks: A bench- marking tool for approximate nearest neighbor algorithms. Information Systems, 87:101374, 2020.
[11] Alexandr Andoni and Daniel Beaglchole. Learning hash robustly, guarantced. In International Conference Machine Learning, pages 599-618. PMIR, 2022.
[12] Joshua Engels, Benjamin Coleman, and Anshumali Shrivastava. Practical near neighbor search via group testing. Advances in Neural Information Processing Systems, 1149950-9967 2021.
[13] Honghao Lin, Tian Luo, and David P. Woodruff. Learning augmented binary search trees, 2022. https http///xxvvv.o///// www.voor//aaa httg/abs//00/ 2206. 12110
[14] Andrej Karpathy. nanogpt. https: http://hhhha..oco/aaaaattttaaaaa 2024. Accessed:
2024-05-28.

---

[15] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
[16] Aditya Grover, Eric Wang. Aaron Zweig, and Stefano Ermon. Stochastic optimization of sorting networks via continuous relaxations. In International Conference on Learning Representations, 2019. URL httpttt://pppereiiieeeeeee
wwt. weet/foorun1ddd-111533303
[17] Marco Cuturi, Olivier Teboul, and Jean-Philippe Vert. Differentiable ranking and sorting using
optimal transport. Advances in neural information processing systems, 32, 2019.
[18] Felix Petersen, Christian Borgelt, Hilde Kuchne, and Oliver Deussen. Monotonic differentiable sorting networks, 2022.
[19] W. W. Peterson. Addressing for random-access storage. IBM Journal of Research and Develop- ment, 1(2):30-146, 1957. doi: 10. 10.1147/rd.12.0111
[20] Jon M Kleinberg. Two algorithms for nearest-neighhor search in high dimensions. In Pro- ceedings of the twenty-ninth annual ACM symposium on Theory of computing, pages 599-608, 1997.
[21] Alexandr Andoni, Piotr Indyk, and Ilya Razenshteyn. Approximate nearest neighbor search in high dimensions. In Proceedings of the International Congress of Mathematicians: Rio de Janeiro 2018, pages 3287-3318. World Scientific, 2018.
[22] Moses S Charikar. Similarity estimation techniques from rounding algorithms. In Proceedings of the thiry-fourth annual ACM symposium on Theory of computing, pages 380-388. 2002.
[23] Chen Yu Hsu, Piotr Indyk, Dina Katabi, and Ali Vakilian. Learning-based frequency estimation algorithms. In International Conference on Learning Representations, 2019. URL https
/openrevieww... weeee/forummidrrlloooooae
[24] Graham Cormode and Shan Muthukrishnan. An improved data stream summary: the count-min sketch and its applications. Journal of Algorithmss 55(1):58-75, 2005.
[25] CAIDA. Caida internet traces, httpecchhhcaoohhhtt////////eeeee http://www.aaaaeeeddaaaa/o/oooos chicago.xml, 2016.
[26] Mikkel Thorup and Uri Zwick. Approximate distance oracles. Journal of the ACM (JACM), 52 (1):1-24, 2005.
[27] Aydın Buluç, John Gilbert, and Viral B Shah. Implementing sparse matrices for graph algorithms. Graph Algorithms in the Language of Linear Algebra, pages 287-313. 2011.
[28] Diptarka Chakraborty, Lior Kamma, and Kasper Green Larsen. Tight cell probe bounds for succinct boolean matrix-vector multiplication. In Proceedings of the 50th Annual ACM SIGACT Symposium on Theory of Computing, pages 1297-1306, 2018
[29] Guy Blanc, Jane Lange, Mingda Qiao, and Li- Yang Tan. Decision tree heuristics can fail, even in the smoothed setting. arXiv preprini
2021.
[30] Guy Blane, Jane Lange, Mingda Qiao, and Li- Yang Tan. Properly learning decision trees in almost polynomial time. Journal of the ACM, 69(6):1199, 2022.
[31] Tim Kraska, Alex Beutel, Ed H Chi, Jeffrey Dean, and Ncoklis Polyzotis. The case for learned index structures. In Proceedings of the 2018 international conference on management of data, pages 489-504, 2018.
[32] Sepanta Zeighami and Cyrus Shahabi. On distribution dependent sub-logarithmic query time of learned indexing. In International Conference on Machine Learning, pages 40669-40680. PMLR, 2023.
[33] Yihe Dong, Piotr Indyk, Ilya Razenshteyn, and Tal Wagner. Learning space partitions for nearest neighbor search. arXiv preprint arXiv:1901.088 2019.
12

---

[34] Jun Wang, Wei Liu, Sanjiv Kumar, and Shih-Fu Chang. Learning to hash for indexing big data-a survey, Proceedings of the IEEE, 104(1):34-57, 2015.
[35] Ibrahim Sabek, Kapil Vaidya, Dominik Horn, Andreas Kipf, Michael Mitzenmacher, and Tim Kraska. Can learned models replace hash functions? Proceedings of the VLDB Endowment, 16
(3), 2022.
[36] Lawrence Cayton and Sanjoy Dasgupta. A learning framework for nearest neighbor search. Advances in Neural Information Processing Systems, 20, 2007.
[37] Anders Aamand, Justin Chen, Huy Nguyen, Sandeep Silwal, and Ali Vakilian. Improved frequency estimation algorithms with and without predictions. Advances in Neural Information Processing Systems, 36: 36:14387-14399 2023.
[38] Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXiv:1410.5401, 2014.
[39] Petar Veličković, Rex Ying, Matilde Padovano, Raia Hadsell, and Charles Blundell. Neural execution of graph algorithms. arXiv preprint arXiv:1910. 2019.
[40] Petar Veličković, Adrià Puigdomenech Badia, David Budden, Razvan Pascanu, Andrea Ban- ino, Misha Dashevskiy, Raia Hadsell, and Charles Blundell. The clrs algorithmic reasoning benchmark. In International Conference on Machine Learning, pages 22084-22102. PMLR, 2022.
[41] Beatrice Bevilacqua, Kyriacos Nikiforou, Borja Ibarz, Ioana Bica, Michela Paganini, Charles Blundell, Jovana Mitrovic, and Petar Veličković. Neural algorithmic reasoning with causal regularisation. In International Conference Machine Learning, pages 2272-2288. PMLR. 2023.
[42] Gleb Rodionov and Liudmila Prokhorenkova. Neural algorithmic reasoning without inter- mediate supervision. Advances in Neural Information Processing Systems, 36:51663-51674, 2023.
[43] Yukun Cao, Yuan Feng, and Xike Xie. Meta-sketch: neural data structure for estimating item frequencies of data streams. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 6916-6924, 2023.
[44] Yukun Cao, Yuan Feng. Hairu Wang. Xike Xie, and S Kevin Zhou. Learning to sketch: A neural approach to item frequency estimation in streaming data. IEEE Transactions Pattern Analysis and Machine Intelligence, 46(1 46(11):7136-755) 2024.
[45] Yuan Feng. Yukun Cao, Hairu Wang. Xike Xie, and Kevin Zhou. Lego sketch: A scalable memory-augmented neural network for sketching data streams. arXiv preprint wrXiv:2505.19561 2025.
[46] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 1607.0660 2016.
[47] Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th international conference on machine learning (ICML-10), pages 807-814, 2010.
[48] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
[49] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
[50] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmat 2017. httph:tttt///ttte htttpa///arxiv.oor//a/////111111
[51] Sunil Arya, David M Mount, Nathan S Netanyahu, Ruth Silverman, and Angela Y Wu. An optimal algorithm for approximate nearest neighbor searching fixed dimensions. Journal of the ACM (JACM), 8916)39 1998.
13

---

[52] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers, 2020.
[53] Irwan Bello, Hieu Pham, Quoe V Le, Mohammad Norouzi, and Samy Bengio. Neural combina- torial optimization with reinforcement learning. arXiv preprint arXiv:1611 161.0940 2016.
[54] Barret Zoph and Quoe V Le. Neural architecture search with reinforcement learning. arXiv preprint arXiv: 161.0577 2016.
[55] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344-16359, 2022.
[56] Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. arXiv preprint arXiv:20/2.09699, 2020.
[57] Chengxuan Ying et al. Do transformers really perform badly for graph representation? In Advances in Neural Information Processing Systems, volume 34, pages 28877-28888, 2021.
[58] Hengshuang Zhao et al. Point transformer. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 16259-16268, 2021.
[59] Alexey Dosovitskiy al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1192) 2020.
[60] Andrew Jaegle, Felix Gimeno, Andrew Brock, Oriol Vinyals, Andrew Zisserman, and João Carreira. Perceiver: General perception with iterative attention. In International Conference on Machine Learning, pages 4651 PMLR, 2021.
[61] Jean- -an--Batistt Borgeaud, Jean- Alayrac, Carl Doersch, Catalin Ioneseu, David Ding, Skanda Koppula, et al. Perceiver io: A general architecture for structured inputs & outputs. arXiv preprint arXiv:2107.14799 2021.
14

---

Appendix
Related Work
17
Nearest Neighbor Experiments
17
B.1
Training Details
17
B.2
Training Variance
18
B.3
Noise Injection for Lookup Sparsity
18
4
Additional ID/2D/30D Uniform Distribution Performance Plots
19
B.
Uniform Distribution
19
B.6
Hard Distribution
20
B.7
ID Zipfian Experiment
24
................................
B.8
Hashing Baselines
24
B.8.1
Locality Sensitive Hashing (LSH)
24
B.8.2 Learning to Hash
24
................................
B.8.3 Baseline Choices for MNIST Experiment
25
B.9 N-8, M=1 30D LSH Experiment
25
B. 10 Additional High-Dimensional NN Experiments with Realistic Data
25
B. MNIST (3-digit) Experiments
27
12 Extra Space
28
B.12.1 ID Extra Space
28
B. 1.12
Bucket Baseline
29
B.1..3 Understanding Extra Space Usage
29
B.12.4 30D Extra Space
30
Frequency Estimation
30
CountMinSketch
30
Frequency Estimation Experiment Details
Improved Frequency Estimation with Augmented CountMinSketch
31
C.4
Learning Heavy Hitter Features]
32
Relation of our experiments to Hsu et al.
33
C.6
Connection to neural sketching works [43 45
33
Frequency Estimation Failure Mode
34
D Efficiency & Scaling
34
D.I Improving Efficiency of Data-Processing Network for NN Search
34
D.2 Scaling Up
34
D.3
Benchmarking Computational Overhead
36
Broader Applicability of Our Framework
37
E.1 Generality of Our Framework
37
15

---

E.2 Nearest Neighbor Permutation
38
16

---

A Related Work
Algorithms Recent work has shown that traditional data structures and algorithms can be made more efficient by learning properties of the underlying data distribution. [1 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 start the search for a record, which can lead to provable improvements to the query time over non-learned structures [32]. Other works augment the data structure with predictions instead of the query algorithm. For example, [8] use learned frequency estimation oracles to estimate the priority in which elements should be stored in treap. Perhaps more relevant to the theme of our work is [33] which trains neural networks learn partitioning of the space for efficient nearest neighbor search using locality sensitive hashing, and the body of work on learned hash functions [4 [35] [T] There also work on learning the parameters of nearest neighbor search algorithms using machine learning [36]. There has also been a number of works that have explored Baariing-augment algorithms for frequency estimation [23 37 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. Our approach eliminates the human-in-the-loop. making promising for settings with limited insights into suitable data structures. However, this comes at the cost of losing the provable guarantees that marning--ugmenttn methods typically offer.
Neural Algorithmic Learners There is a significant body of work on encoding algorithms into deep networks. Graves et al. (38] introduced the Neural Turing Machine (NTM), which uses external memory to learn tasks like sorting and copying. Veličković et al. [39] used graph neural networks (GNNs) to encode classical algorithms such as breadth-first search. These works train deep networks with great degree of supervision with the aim of encoding known algorithms. For instance, Graves et al. [38 use the ground truth sorted list 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 al. 2 train neural networks to solve SAT problems. Garg et al. B 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.
B Nearest Neighbor Experiments
B.1 Training Details
The transformer in the data-pr processing network is based on the NanoGPT architecture 14 and has 8 layers with 8 heads each and an embedding size of 64. Instead of using discrete tokens, we have input and output projection layers (Linear layers) that project the input to the hidden dimension and from the hidden dimension to the scalar output values. Each query model Q is 3-layer MLP with a hidden dimension of size 1024. Each hidden layer consists of a linear mapping followed by LayerNorm (46) and the ReLU activation function [4]. In all experiments we use a batch size of 1024, 1e-3 weight decay and the Adam optimizer [48] with default Py [49] settings. We do a grid search over {0.0001, 0. 0.00001, 0. 0.00005} to find the best learning rate for both models. Near the end of training when performance plateaus, we decrease learning rates to boost performance. We apply the Gumbel Softmax (50) with temperature of to the lookup vectors to encourage sparsity. Apart from the learning rate of the two models, we did very minimal hyperparameter tuning as the majority of experiments worked without additional tuning. We expect additional performance gains would come with further tuning but are more interested in understanding the general performance landscape relative to algorithmic baselines rather than optimizing any single result. All experiments are run on a single NVIDIA RTX8000 GPU.
Training Steps/Data set Sizes As we have direct access to the dataset distribution, for all ex- periments (except those with real high-dimensional data), we do not use a fixed training set but rather train 'online", sampling new (training) datasets directly from the dataset distribution. The number of optimization steps we use varies across experiments but is up to 500k for most experiments with N < 100, though for experiments with N 500 we ran for up to 2 million steps (with carly
17

---

stopping). For experiments with fixed training sets (e.g. 3-Digit ANISTFFasiionMNIST/SIIITT we resampled from their corresponding training sets (sizes 60k/60k/1M respectively) and used their test sets for evaluation. Given that the models perform well on test data for these distributions, it suggests that we do need much unique data as we use the online setting. We leave an exploration for this to future work.
B.2 Training Variance
We found that the setup is robust to random model seeds, however, below we include plots of ID/2D/30D experiments using three different seeds to showcase the variation.
Model seed
Model seed
0.35
Model
seed
0.6
2000
0.30
Socd
0.6
Accuracy
0.25
0.4
0.3
0.2
0.10
0.0
0.1
0.05
Lookupp Index
Lookupp Indox
Lookup Index
Figure 6: ID (Uniform), 2D (Uniform), 30D (Hypersphere) experiments with different seeds.
B.3
Noise Injection for Lookup Sparsity
We find that adding noise prior to applying the soft-max on the lookup vector m, leads to sparser queries. We hypothesize that this because the noise injection forces the model to learn noise-robust solution, and the softmax output becomes robust when the logits are well-separated. Well-separated logits, in turn, lead to sparser solutions.
Consider simplified setup in ID where the query model is not conditioned on g and is only allowed
one lookup (M = 1) and /) is sorted list of three elements: D = (2, x33] For given query
and its nearest neighbor y. the query-execution network is trying to find the optimal vector ni e R [50]. Given
that minimizes | m DI2 where m softmar((m e), e Gumbel distribution
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 my mig mag However, when noise is added to the logits in this solution is destabilized. Instead, in the presence of noise, the model can robustly select the middle element by making my much greater than mn my. 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 N/2-
18

---

B.4 Additional ID/2D/30D Uniform Distribution Performance Plots
Model
0,7
E2E
(fazen)
(nen-adaptiva)
0.5
KD Tree
0.3
0.1
0.0
Lookup
Figure 7: 2D Uniform Accuracy.
Model
Model
Square Error
EZE (non adaptive)
Square Error
0.99
EZE (non adaptive!
KD Tree
0.20
LSH (K+ 3)
Square
E2E
Itrazeni
20
adaptive)
E2E Ino permute!
Binary Search
6.m
Lookup Index
Lookup Index
Lookup Index
Figure 8: Mean square error plots for (Left) ID Uniform distribution, (Center) 2D Uniform distribu- tion, (Right) 30D Uniform distribution over unit hyper-sphere.
B.5 2D Uniform Distribution
Original position cleseness
by position cleseness
0.75
0.75
0.53
0.50
0.23
0.25
8.6
8.4
-1.03
-1.00
Dinodiin
Figure 9: Our model's learned permutation on the 2D uniform distribution. The model puts elements that are close together in the Euclidean plane next each other in the permutation.
19

---

330
van
142
-222
-022
-1.00
-2.5
403
1.30
-303
Modil
DAT
CNIY
530
225
-3.2-
-023
-032
-1.00
-352
1.00
0.7
330
3.25
300
336
8.75
0.3
0.53
0.25
C50
Due Mudel
Cany
Oeery
90
60
322
-0.25
-035
-0.00
-052
+0.0
-0.2
13.05 1275
6557 022 800
325
330
100
1.00
0.75
0.50
0.25
032
625
030 073 300
Figure 10: 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.
B.6 Hard Distribution
To generate data from the hard distribution, 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 smaller range and The intuition behind this distribution is to reduce concentration such that p( Nlg) is roughly uniform where NN denotes the index of the nearest-neighhee of in the sorted list.
Precisely, to sample points from the hard distribution we generate a random balanced binary tree of size N. All vertices are random variables of the form Uni form(0 n-k) where is some constant and is the level in the tree that the vertice belongs to. If the th node in the tree the left-child of its parent, we generate the point 2; H; Tp(x) d, where p(i) denotes the parent of the th node and d, is a sample from node of the random binary tree. Similarly, node the right child of its parent, I; p((s) + dz. For the root element Zg dg. In our experiments we set concenennnt...on. We found it challenging to train models with N 16 as the range of values that F; take increases with N. Thus for larger N, the model needs to deal with numbers at several scales, making learning challenging.
a = 7. The larger the value of a, the greater the degree
20

---

1.0
Model
E2E
Binary Search
Accuracy
0.6
0.2
Lookup Index
Figure 1I: Our model's performance is closely aligned with binary search on the hard distribution in ID. By design, this distribution does not have useful prior our model can exploit and it learns binary search like solution.
10 Féenibu on Average molip Bextion
Figure 12: The positional distribution per lookup in the 1D Hard experiment. Our model closely aligns with binary search, first looking the middle element, then (approximately) either the 25th or 75th percentile elements, and so on.
21

---

Binary-Sarch Scarch
Our Model
Lookup tridet
Binary
1an
X
21
Loceup Inden
Lookup Indee
Binary
Jookup index
Lookup Indox
Modal
Indan
Irides
Figure 13: Binary Search vs. our model on the hard distribution in ID.
22

---

1.0
Model
E2E
Tree
Accuracy
0.6
0.2
Lookup Index
Figure 14: On the 2D hard distribution model roughly tracks the performance of k-d tree.
23

---

B.7 ID Zipfian Experiment
1.0
Accuracy
0.4
Model
0.2
Binary
Zipf Treap
Lookup Index
Figure 15: For 1D Zipfian query distribution, our model performs slightly better than the the treap algorithm from [3] and both methods significantly outperforms binary
search.
Prior work has shown that several real-world query distributions follow Zipfian trend whereby few elements are queried far more frequently than others, leading to the development of learning- augmented algorithms aimed at exploiting this [33. We consider a setting where PD is the discrete uniform distribution over {1,..., 200} and Pg is Zipfian distribution over {1, 200} skewed towards smaller numbers such that the number i is sampled with probability proportional to # We
set a = 1.2. Again, in this setting N = 100 and M 7.
In Figure we compare our model to both binary search and the learning- -uggmented treap from Lin al. 8 Our model performs slightly better than the learning-augnin -uggennd treap and both algorithms significantly outperform binary search with less than (Ng(N) queries. This setting highlights crucial difference in spirit between our work and much of the existing work on saarinn--uugment 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-loog This is promising could be useful in settings where we lack insight on appropriate data structures. The flip side, however, is that barning--aumenter data structures usually come with provable guarantees which are difficult to get when training models an end-to-end fashion.
B.8
Hashing Baselines
B.8.1
Locality Sensitive Hashing (LSH)
Our LSH baseline is based the SimHash algorithm by 22 SimHash is particularly relevant for our hypersphere setting as the algorithm is designed for nearest- neirest retrieval based on cosine similarity and all the points in experiment fall on the unit hypersphere so minimizing euclidean distance is equivalent to minimizing the cosine similarity. We construct the LSH baseline as follows. We sample K random vectors rı, ...; FK from the standard normal distribution in Rdd For given vector E IRd, its hash code computed as hash(v) = [sign(v (onn(vrr))) (rrr ),.., sign(v (v (ogn(vrr))-- In total, there are 2 possible hash codes. To create a hash table, we assign each hash code bucket of size N/2K, For given dataset D = (a.,...,,II
we place each input in its corresponding bucket (determined by its hash code hash(x; (x,). ). If the bucket is full, we place 2; 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(g). If there less than M vectors in the bucket, 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). Note that for all experiments that compare to LSH, we choose the value of K that maximizes the performance of LSH.
choose
B.8.2 Learning Hash
In addition to locality-sensitive hashing, we include comparisons to several learning-tor thsh baselines: ITQ, K-Means, and NeuralLSH. See Dong al. [33] for more details about these specific basclines.
24

---

Similar to classical LSH methods, learning- to-hash methods aim to hash elements to buckets such that close by elements are more likely to be mapped to the same buckets. The main difference is that learning to hash methods learn the hash function from data while it is traditionally hand-designed. However, note that while our framework learns the whole data structure and an adaptive query algorithm in an end to end manner, the learning to hash methods just learn part of the data structure while the query algorithm is the same as that of LSH and thus is not adaptive.
B.8.3 Baseline Choices for MNIST Experiment
For the 3-digit MNIST setting, we exclude the learning-to soalinng-t-hhasttt hoshhshh methods as they cannot be applied directly. Specifically, in our MNIST setup the metrie is not given but implicitly provided via nearest neighbor supervision. This means our model needs learn the metric space, which learning-to-bi to-hahhh methods are not designed for. However, to highlight our models' strengths in this regime, we compare with an oracle hash function (denoted as LSH in the table)- splitting the 200 number universe evenly into K clusters where the first N/K elements are assigned to cluster 1, the second N/K elements to cluster 2 and so forth. Note that no learning-to-har tosss method could outperform this oracle hash function it is upper-bound with access to the correct metric. To execute a NN search, query mapped to its corresponding bucket and the element closest to this number (defined over this ID label space) is chosen the NN. When N=50 and M=7 we find that the optimal K is 8.
In Figure 1y we also compare with binary search directly operating on the digit space. To be clear, these are unfair comparisons as our model has to learn features that capture the underlying metric space but in this case is provided directly for the two baselines. The oracle hashing baseline (and thus any learning-to ho-hash baseline) should still underperform our model in this regime as it does not use an adaptive query algorithm which is necessary for achieving the (log.N) performance our model recovers.
B.9 N=8, M=1 30D LSH 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 's mapping can be explained by the projection onto these two components. In Figure 16 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 E D and color them by the position they were permuted to. The plot on the right suggests that the data-proceet paot--prcessing 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 lookup in the position that corresponds to the query vector's 's bucket. This behaviour suggests the model has learned locality-sensiteee hashing type solution! Specifically, both the data- data-processio processing and query models learn the same partitioning of space in tandem. These partitions serve the same role as hash functions in LSH algorithms like SimHash [22]. Like LSH algorithms, to find neirest-, the query model maps point to its corresponding hash bucket and then compares it to other points in that bucket, ultimately selecting the closest one.
B.10 Additional Hligh-Dimensional NN Experiments with Realistic Data
We run additional nearest-neighbor experiments on two standard high- dite-iinnniio approximate nearest-neighhor benchmarks: SIFT and FashionMNIST (projected to 100 dimensions via PCA) [10 to further demonstrate our model can handle more realistic data. The SIFT dataset contains high-dimensional feature descriptors extracted using the Scale-Invariant Feature Transform (SIFT) algorithm. It is commonly used for nearest neighbor search, image retrieval, and feature matching tasks. The FashionMNIST dataset consists of 28x28 grayscale images of clothing items, such as shoes, shirts, and dresses.
25

---

Query Projectin onto top 2 PCA components
Permuted Dafaset Projection anto tap 2 PCA tompanents
0.0
00-
LEOAUD HoOioon
POTTREDO POODEn
-42
-0.4-
-06
-000
-0. 0. 58
-0.6
-0.4
-0.2
2.0
0.6
Figure 16: (I.eft) Projection of queries onto two PCA components of the decision boundaries of the query model, colored by the lookup position the query is mapped to. (Right) Projection inputs onto the same PCA components colored by the position the data-processing model places them in. Both the pate-proessiing and query models map similar regions to the same positions, suggesting an LSH-like bucketing solution has been learned.
MNISST (V- IC0, M-6, d-100)
SIFT (N-100, M-6, d-128)
0.7
Aleptiie
Adaptive
LSI
0.6
Nourol.SH Hash
NeuralSSH Has?
0.4
0.3
0.1
0.1
Lockup ndox
.ockuup indcx
Figure 17: Our model compared to various hashing baselines on the FashionMNIST dataset (left) and the SIFT dataset (right). At the final query, our E2E solution performs as well as the hashing baselines. At intermediate lookups, it slightly under-performs, however, the model is only being trained to optimize performance the final lookup. We not expect our model to outperform these baselines in this setting as it is unclear that query adaptivity should be beneficial here. Rather, we include these results further emphasize that our E2E model can recover reasonable solutions in variety of settings even when compared to carefully hand-designed solutions.
For each dataset, we use the train split to train our model (as well as the dearning-to-hash baselines) and for evaluation we use the test split. The learning-to-Be taaingg-t-hhsh basclines (described in B. 8.2.) cluster the data into 16 partitions and then the same setup described in Appendix B.8. 1 to execute the NN search. chose 16 partitions as this produces the best performance for the learning-to- baselines in setup where N 100 and M 6.
Our model performs competitively with these learning serninggtt-hhaee to-hash baselines (Figure 17p. expect our model to outperform these baselines in this setting as it is unclear that query adaptivity should be beneficial here. Rather, we include these results to further emphasize that end-to-end learning can recover reasonable solutions in variety of settings even when compared to carefully hand-designed solutions. We also show an example of how the data-procsssing model learns to transform the FashionMNIST data by organizing the images by class (Figure 18).
We do not
26

---

Figure 18: (Left) Sample dataset from FashionMNIST (Right) The learned data structure. The data-processing model learns to cluster similar items together.
1.0
Model
E2E
0.8
Binary Search*
LSH*
Accuracy
0.6
0.4
0.2
0.0
1
2
5
6
7
Lookup Index
Figure 19: 3-Digit MNIST Nearest Neighbors Accuracy. Even though binary search (over the underlying digits) is an unfair comparison, we include it reference to compare our model's performance with.
B.11 MNIST (3-digit) Experiments
We generate images of 3-digit numbers by concatenating digits from MNIST (see Fig. 20 for image samples). To construct nearest-neighbor dataset D, we sample N 50 labels (each label corresponds to a number) uniformly from to 199. For each label, we then sample one of its associated training images from 3-digit MNIST. Additionally, we sample query label (uniformly over (0,., .., 199}) and its corresponding training image and find its nearest neighbor D, which corresponds to the image with the closest label. We emphasize that the model has no label supervision but rather only has access to the query's '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- procespnccssing and query-execution networks should operate over the same low- dimensional representation we train a CNN feature model (architecture described below) Fg as well. Our setup remains the same as before except now the data-processing network and query- execution network operate on (F1( ..., (x(x)) and Fa(g). respectively. As the underlying distance metric does not correspond the Fuclidean distance, we minimize the cross-entropy loss instead of the MSE loss. Note that the cross-entropy loss only requires supervision about the nearest
27

---

0
6
0
7
1
0
9
/
8
3
Figure 20: Samples from 3-Digit MNIST
Model
2
8
16
32
64
128
Bucket Baseline
5.9
11.7
23.6
44.5
65.8
81.1
89.5
E2E
44.0
47.9
52.0
58.8
71.9
75.3
88.4
91.3
Table 2: ID NN search accuracy given varying amounts of extra space. The E2E model effectively uses extra space and outperforms a bucketing baseline.
neighbor of the query, and does not require the exact metric structure, so it can used even where the exact metric structure is unknown.
CNN Feature Model Architecture F6
(inncinnnnnc out_channels=32, kernel_size=3, stride=1, padding=1) ReLU()
MaxPoo126 (kernel_size-2 2 padding
Conv2d (in_iiinneeeel______ 32, out_cchaneess64 kernel_size= stride=1, padding-1) ReLU()
MaxPoo12d (kernel stride=2, padding=0)
Linear (in_ faln_ffeatures=6447221111 out ant_featuues=1128 ReLU()
Linear fin_features=128, out out_features==)
We plot the results in Fig 19 compared to optimal hashing and binary search baseline. Note these are not fair comparisons as these baselines operate over numbers directly instead of their corresponding images and are only provided for comparison. Please refer to B.8.3 more detailed discussion on this.
B.12 Extra Space
We also consider scenarios where the data structure can use additional space. To support this use case, the drocesssn transformer outputs T extra vectors b ..., by E Rd which can be retrieved by the query-execution network in the same way as the other outputs. We form the data structure ò by concatenating the permuted inputs and the extra vectors: ò = [Dp, br.,
B.12.1 ID Extra Space
Most of our nearest neighbor experiments show our model's ability to learn useful orderings for efficient querying. Data structures, however, can also leverage pre-computed information to accelerate search. For example, with infinite space, data structure could store the nearest neighbor for every
28

---

possible query, enabling ((1) search. Here, we test whether the model can effectively use additional space.
We run an experiment where the data and query distribution are uniform over (-1.1) with N =
50, M 2. We allow the prota-ppeeesingg network output T (0,21, 22, 23, 2 25, 26 2 numbers b.., bt E R in addition to the N rankings, We plot the NN accuracy as function of 'I' in Table (aaleeecmmrreed
simple bucketing baseline (explained
B.1.2.
Our model's accuracy monotonically increases with extra space demonstrating that the data-procsssing network learns pre-compute useful statisties that enable efficient querying. We provide some insights into the learned solution in App B. 12.1 and show that our model can be trained to use extra space in the nigh-dimensional case well (App B.12.4.
Average Lookup Posibion
04
50
Extra space position:
Extra space position:
space position: 37
Bigressn on Coetrdaer
Creficeer
8.6
Bigita on Crefddar
0.
0.2
02
(Sorted) trput Peaition
(Sorted
(Sorted) Irpus Postion
Extra position:
1.00
Way cavoen Goifiios
0.7
0.50
0.20
5:25
0.25
2.00
10
20
10
20
10
2-
(Samedi Irpuet Ponnion
(Sorted) Ingurt Rovdit
[Sorted Irpurt
Figure 21: (Top) Decision boundary of the first query model. (Bottom) The regression coefficients
of the values stored extra positions linear function of the (sorted) inputs.
B.12.2 Bucket Baseline
We create simple bucket bascline that partitions [-1, 1] into T evenly sized buckets. In each bucket b; we store ar argminz,e (egnnnnenllls..... e - 4|l where 4 is the midpoint of the segment partitioned in b, This baseline maps a query to corresponding bucket and predicts the input stored in that bucket as the nearest-neighbor. As T - oo this becomes optimal hashing- like solution.
B.12.3 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 21 we plot the decision boundary of the first query model. The plot demonstrates that the model chunks the query space (I-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 values stored in each of the extra space tokens b, and plot the coefficients for several of the extra spaces in Figure 21 For a given subset of the query range, the value stored at its corresponding extra space approximately a weighted sum the values stored the indices that correspond the percentile of that query range subset. This useful information as it tells the model for given query
29

---

percentile how 'shifted" the values in the current dataset stored in the corresponding indices are from model's prior.
B.12.4 30D Extra Space
0.6
Model
E2E
0.5
E2E (Caefficients)
Baseline
Accuracy
0.4
0.1
2
2
Extra Spaces
Figure 22: Our unconstrained model (E2E) and more interpretable version (E2E (Coefficients))
both learn to effectively leverage an increasing amount of extra space in 30D, with the unconstrained model outperforming an LSS baseline.
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 [5T].
We follow a similar setup to the ID 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 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 demonstrate the versatility of our framework. If one is only interested in identifying strong lower-bound of how well one can use 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 22 While both models perform competitively with LSH baseline and can effectively leverage an increasing amount of extra space, the unrestricted model outperforms the coefficient model at a certain point.
C Frequency Estimation
C.1 CountMinSketch
CountMinnSetch [24] is probabilistic data structure used for estimating the frequency of items in a data stream with sublinear space. uses a two-dimensional array of counters and multiple independent hash functions to map each item to several buckets. When new item x arrives, the algorithm computes d hash functions h (1)x, h2(x),.... h,(x), each of which maps the item to one of buckets in different rows of the array. The counters in the corresponding buckets are incremented by 1. To estimate the frequency an item æ, the minimum value across all counters C[1, (=)), C(2, hz(z)), ... Cid, hz(z)] is returned. The sketch guarantees that the estimated frequency f(x) of an item x is at least its true frequency f(x), and most f(x) + eNV, where N is the total number of items processed,
L. and w is the width of the sketch. The probability that the estimate exceeds this bound is at most 5 1. where d is the depth of the sketch (i.e., the number of hash functions). These guarantees hold even the presence of hash collisions, providing strong worst-case accuracy with O(w d) space.
30

---

C.2 Frequency Estimation Experiment Details
Here, we outline our approach to modeling the frequency estimation problem. Following NN search, we chose this problem to further explore the broader applicability of the end-to-end learning paradigm for the following reason: its streaming nature makes fundamentally different from NN search, requiring adaptations to the framework used in NN search. In contrast, the other problems discussed in Section 3.2 require minimal adaptations (see App E That said, similar to NN search, the two key constraints in this problem remain the size of the data structure and the number of lookups. As a result, the high-level setup remains the same: we still use a sata-processing network and query network, applying similar principles to optimize data structure efficiiecy. For instance, control space complexity by tokenizing and restricting the data structure's 's size and enforce query complexity using similar sparsity techniques networks. Note that while there are differences between the architectures used for NN search and frequency estimation, we explain how both problems fit into the broader framework in section B
lookup vectors. Next, describe the aata-processing and query
Data processing Network We model the data structure as a k dimensional vector Đ and use an MLP as the pata-procsssing network which responsible for writing to Đ. 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 cata-processing network outputs M I-dimensional one-hot update position vectors uI...., M and M corresponding scalar update values UI...., UM- We then apply the update, obtaining D D, M constrain the construction complexity of the data structure, here we have limited each update to the data structure budget of M lookups. We do
=i=1= U, U.- Unlike in the NN setting where we did not
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 handled in similar fashion to NN search we have M query MLP models that output lookup positions. Finally, we also train MLP predictor network (V1,,.., UM) that takes in the M values retrieved from the lookups and outputs the final prediction.
Training Details We follow the same setup as the nearest neighbors training except for frequency estimation, the data-pnn mata-processing network is 3-layer MLP with hidden dimension of size 1024. We do a grid search over {0.0001, 0.00005, 0.000013 to find the best learning rate for both models. Models are trained for 200k gradient steps with early stopping. All experiments NVIDIA RTX8000 GPU.
on single
C.3 Improved Frequency Estimation with Augmented CountMinSketch
In our experiments on learning frequency estimation algorithms (Section 3.1 we found that on the Zipfian distribution our model was able to outperform the CountMinSketch algorithm by using smaller update delta. 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 least partially due the smaller delta.
We use this insight to design a modified version of CountMinSketch that uses a custom update delta. In Figure 23 we show that this augmented CountMinSketch algorithm outperform vanilla CountMinSketch on the large-scale CAIDA IP traffic dataset (25) by up to a factor of two. These results demonstrate that even small scale, our model can provide useful insight into data structure design that can be transferred to realistic settings.
The CAIDA dataset (25 consists of traffic data collected in 2016 from a backbone link of a Tier- I ISP between Chicago and Seattle. Each recording session spans approximately one hour, capturing around 30 million packets and I million unique flows per minute. We use the first minute for our experiment.
Recent work by Aamand et al. [37 independently arrived at similar insight to the one our model discovered automatically. They propose a modified version of CountSketch that improves over standard
on Zipfian distributions, even without explicit predictions- relying only on knowledge of distributional skew. In comparison, our trained model learns to lower the update value uniformly for all clements, strategy that mirrors the motivation behind their use
31

---

of conservative updates, where frequency estimates for some elements are set to zero. A theoretical comparison between the two approaches, as well as a deeper analysis of the learned update rule, is an interesting direction for future work. More broadly, the alignment between our automatically discovered strategy and their hand-crafted variant further supports the idea that small-scale learned models can yield transferable insights with potential for provable guarantees.
Update Della
Best Delta Err/Delta-1 Err
0.7
Absolute Error
0.8
Memory Size
Memory Size
Figure 23: (Left) CountMinSSetch performance the CAIDA dataset with different update deltas vs. memory size. (Right) The relative performance of the best update delta vs the default delta (A 1) for different memory sizes. In some regimes, our augmented CountMinSketch can perform up to twice well as vanilla CountMinSketch just by modifying the update delta.
C.4 Learning Heavy Hitter Features
Model
Mean Absolute Error
5
E2E
CountMinSketch
4
3
2
0
40
60
80
100
Stream Index
Figure 24: On the MNIST heavy-hitters frequency estimation experiment, model significantly outperforms CountMinSketch. This is because our model can learn features predictive of heavy hitters. opposed to the catributioo-agnostic CountMinSketch.
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. (33] show that in frequency estimation for network flows, certain types of IP addresses receive much more traffic than others. We simulate such setting by fixing the rank ordering of the Zipfian distribution. However, instead using a universe of integer elements (1, ..., K}. we instead use their corresponding 3-digit MNIST images with K = 100 (constructed 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 corresponding image from the training set. During evaluation, use images
32

---

Query Model Learned Mapping
10
20
08
Value
0.6
50
02
Memory Position
We show the decision boundary learned by the query/data-prooessing network in the MNIST heavy hitters experiment. As images with smaller numbers occur more frequently in the stream, the wemooy-constrained model learns to reserve separate memory positions for these items in order to prevent collisions among them.
Figure 25:
samples from the test 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, also use feature-leamingg CNN model process the images before passing them to the data-processing and query-execution networks.
We compare model CountMinSketch with l-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 stream of size 100 (see Fig. 24) the latter is iettributioninndpeendeee data-distribution end-to-end, our framework is able to simultancously learn features of heavy hitters (in this case, clustering images with the same label) and this information design 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. 25)
By training from the
C.5 Relation of our experiments to Hsu al.
Here, we clarify the difference hetween our frequency estimation experiments and the setup explored by [2 The frequency estimation experiment in the main paper (Figure[5[(Lee]) is different from the work [23] as we randomize the rank ordering of elements for each dataset so there are not consistent features of "heavy hitters" (frequently appearing elements). Rather, the models can only leverage the fact that the distribution is skewed (without knowing the frequency of any given element). we include an experiment with MNIST data where the ranking fixed but the model must learn features of heavy-hitters (some numbers appear more often than others). The model learns do something similar to what 23
However, in Appendix C.4
propose using separate buckets store counts
for heavy hitters.
C.6 Connection to neural sketching works 43
45
The recent line of work on neural sketching algorithms with our rrequencyesstmaatin
45 shares several similarities experiments. Both can be regarded as memory-augmented neural networks that learn read from and write differentiable memory (8). Although we did not run direct comparisons, their architectures and training setups-tailored tailored specifically for frequency estimation-suggess performance comparable to, and potentially better than, ours, and likely offer more scalable path for this task. By contrast, our goal is different: rather than building in inductive biases from known sketching algorithms to achieve scalability, we deliberately avoid such priors to ask more fundamental question- neural networks discover sketching algorithms from scratch?
33

---

C.7 Frequency Estimation Failure Mode
In our frequency estimation experiments, we found a setting where learning struggled to make progress. Specifically, this occurs when the query and streaming distributions are Zipfian with large skew. For fixed training instances, the rank order of elements is consistent across both stream and query distributions but randomized across different training instances.
We explain our hypothesis for why this occurs with a simplified setup. Imagine that the universe of elements that the query/streaming data is drawn from only consists of two elements: {A, B). Assume the data structure size = 2, thus the memory is large enough to store the counts for each element in separate buckets and therefore incur zero estimation error. Let M [xr, y/ and mn [u, v] where A denotes the initial lookup distribution for element A and m the distribution for B. In other words, when element A arrives, the memory Ó is read by computing m (and in a similar fashion for B). Now, consider the case when both > y and u > V (this applies also when and u < v). If we sample a batch of datasets, in approximately half of them the frequency of A will be much larger than the frequency of B and in the other half the inverse will be true. Again, this is because while the degree of skew is fixed across datasets, the rank ordering of elements is randomized. Consequently, this means that for data sets where A more frequent than B, the optimizer will try to increase v and lower y and when the inverse is true, the optimizer will try to increase u and lower v. However, this means that for both elements, the lookup positions will point to the first position in memory over the second (the opposite is true when x < y and u < v). This will cause collisions and increase the estimation error. This behavior can be avoided by setting e > u, or more generally, making the lookup distributions more orthogonal.
y
While this is a relatively simplified setup (i.e. we're not using any neural networks), we believe it captures part of the more general phenomenon. It demonstrates that there are settings where careful initialization can be required for learning to make progress. More generally, this suggests that end-to-end learning cannot always be expected to work without additional problem specific inductive biases. our view, it also makes more surprising that for the majority of settings we explored, we did not encounter such issues!
D Efficiency & Scaling
D.1 Improving Efficiency of Data-Processing Network for NN Search
For our NN experiments, we chose to use transformer with O(z²) complexity in order to keep the model relatively general. This may not always be necessary. For instance, high-dimensional hashing-basee based methods could be learned with less complex model. However, data structures like k-d trees have a higher construction complexity Odn log(n)). Our framework does not require using any specific model transformer. One way to enable scaling to larger instances would be to lower the complexity of the data-procsssing model (e.g. using linear transformer).
less (or more) complex model can trivially be substituted in for the
To demonstrate that in some settings quadratic attention can be substituted with a cheaper alternative, we run additional experiments in 2D (uniform distribution) and in high dimensions (FashionMNISTT) as the quadratic attention model and plot the results in Figure 26 The comparable performance of the linear attention model suggests that it could be computationalll--cee chehapee alternative that can enable scaling up models to larger settings.
with the lincar attention Performer model [52] We use the same model
D.2 Scaling Up
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 27 (Left/Center)). though with less than log(NN queries. There are two primary limitations that led to relatively low N: the droc-ssogg cost (since all N data points are processed) and the MLP query models, which output vectors of size Nacausing heavy matrix multiplication and memory costs. While we demonstrate that useful data
We use the Pytorch implementation from http:///shhcccccooocoomm//mooomc
34

---

2D Uniform Distribution (N=100, M=6)
FashionMNIST (N=100, M=6)
0.7
Model
Model
Transformer)
[Quadratic Transformer
0
KD Tree
E2E Ilinear Transformer]
[2E (Unear Transformeri
Accuracy
0.5
Accuracy
0.5
0.3
02
0.1
0.1
Lookup Index
Lookup Index
Figure 26: Performance of the quadratic attention transformer vs. linear attention performer trans- former
12
in both 2D (left) and high dimensions (100) (right).
Modil
0,0600
Model
100
Modci
0.4
E22
E22 Ince-ataelivel
Sintry Search
Erro
LSH (C-5)
0.8
EZE cshored)
0.0550
LSH (CoF)
LSH (S=7)
Accuracy
0.6
0.2
0.0500
0.4
0.0475
0.2
0.0
0.0425
0.0
Lookup Index
Lookup Index
Lookupp Index
Figure 27: We scale both the ID (Left) and 30D (Center) experiments datasets of size 500. (Right) We compare our E2E model with version where the query-execution network is only composed of one query-model (E2E (shared)) that is used in loop for M queries during training on the ID Uniform distribution, thereby conserving parameters by reusing weights. This could be promising direction for problem settings where there is a recursive structure to the query algorithm.
structures can still learned at this scale, possible that other classes of structures only emerge for larger datasets.
One avenue to scale end-to-end learning larger datasets is by increasing the parameter count of the Bata-processing and query-execution networks. Morcover, as transformers become increasingly efficient 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. While helpful, these changes would require resources beyond those available with academic compute.
and Complementary to our work, it could also be valuable to explore better inductive biases for the query networks, and other methods 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 first step in this direction show that single query model can be used in-a-loop for NN search in 1D (Figure 27 (Right)). We leave further investigation for future work.
Another inductive bias that could improve scaling would be some form of hierarchical indexing, For example, the data- processing model could assign data points to partitions in a hierarchy, while the query network outputs a smaller set coordinates to index this hierarchy.
These are just initial ideas and there is a large space for inductive biases worth exploring in future works. We deliberately avoided adding such inductive biases as we wanted to leave the setup relatively general and see if anything useful could be learned even at smaller scales. We also emphasize that many of the insights that can derived from our models' learned solutions would scale to larger N. For instance, k-D trees 2d and locality-sensitte hashing in higher dimensions. We limit ourselves
35

---

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 any reasonable scale.
Finally, we point to other works focused on algorithm learning and discovery that started at small scale with more synthetic setups 2 (33 54 and inspired follow-up work that then focused on scaling up the initial set of ideas. We hope our work can serve a similar purpose demonstrating the value in learning data structures E2E and inspire future work develop methods for better scaling.
D.3
Benchmarking Computational Overhead
In this section we evaluate the throughput of the nearest neighbor models across various settings. We emphasize that in this paper we did not optimize for practical efficiency as we were primarily interested fundamental tradeoff's between space and query complexity, though these are certainly related to runtime in practice. Nevertheless, we include our findings for completeness. We use the same architectures from our experiments (see App [.I for details). Benchmarking is conducted with batch size of 512 and the results averaged over 10 trials on single GPU.
Query Processing Throughput
The table below shows the number of queries processed per second when using log(n) lookups per query. As expected, querying smaller datasets has higher throughput as the number of lookups, and thus neural net forward inferences, is smaller.
Dim
N = 64
N = 128
N = 256
N = 1024
N = 2048
32
212,943
175,559
146,495
113,960
103,476
64
211,518
173,807
149,603
113,249
105,833
128
208,741
172,275
151,533
116,607
104,858
Below, we fix N 256 and measure how increasing the number lookups (M) affects throughput. There is a clear inverse correlation between M and throughput indicating that query complexity has implications for practical efficiency.
Dim
M=1
M=2
M=3
M=4
M=5
M=
M=7
M=8
32
1,072,476
517,067
345,946
261,705
212,202
176,503
145,927
134.165
64
1,090,290
513,129
355,654
265,753
211.991
174,031
150,917
132,594
128
1,054,366
519,691
348.347
260,985
212,237
173,406
145,191
133,278
Data-Preprocessing Throughput
The table below shows the number of raw datasets processed into data structure per sec- ond using quadratic and linear attention. Quadratic attention runs faster here due to flash-attention [55], which optimizes memory access. Notably, in smaller N regimes, linear attention can be less efficient than quadratic attention with hardware optimizations, despite its lower FLOP count.
Arch
Dim
= 64
128
= 256
N
1024
2048
Linear
32
17,701
8,772
4,542
1,170
585
Linear
64
17,241
8,844
4,563
1,177
590
Linear
128
17,121
8,614
4,473
1,152
577
Quadratic
32
95,229
49,697
13,418
5,867
571
Quadratic
64
94,815
46,924
13,285
5,248
548
Quadratic
128
93.294
48,506
13,095
5,174
543
We also try fixing the input dimension to 128 and adjust the number of layers the quadratic transformer.
36

---

# Layers
64
N 128
N 256
N 1024
2048
174,433
174,381
173,501
31,512
8,449
8
94,493
95,037
18,849
6,780
610
12
65,331
24,301
8,073
1,103
220
E
Broader Applicability of Our Framework
E.1 Generality of Our Framework
Our framework designed for problems that require efficiently answering queries over some col- lection of data. We believe it can be applied to any problem that shares this structure-named namely, one where there data-processing step, a query algorithm, and constraints on space and query-time complexity. In Section 3.2 we outline several such candidate problems that could benefit from this approach. Many classical data structure problems share these common elements, and our framework provides a unified way of modeling them.
key feature of this formulation is the enforcement of query-time complexity constraints, which requires learning sparse querying algorithms. We address this by training with noise to promote sparsity (see App. B.3 a core mechanism that makes the framework practical in constrained settings.
In its most general form, the framework consists three components: data-processing network that maps N data points into K structured chunks; query model that performs M lookups over these chunks based on a given query; and a predictor network that produces the final response.
In nearest neighbor search, we have N = K, and the inputs must be retained during data structure construction. This motivates learning permutations as part of the data model. As discussed in App. E.2 this approach does not limit generality and can apply any setting where inputs must be preserved. Additionally, because the final response is itself a data point, no separate predictor network is needed.
In contrast, frequency estimation introduces structurally different challenge: N K (stream length vs. memory size), and the data-processing network uses an MLP instead of a Transformer to better suit the streaming setting. Here, we also include predictor network to estimate frequencies, while reusing the same query mechanism from the NN setup. We chose frequency estimation our second problem precisely because it structurally different due to the streaming nature, yet E2E learning still works under the proposed framework.
Adapting Framework for Other Problems
As nearest neighbour search is arguably the most pervasive data structure problem, there remain many other avenues to explore end-to-end learning for nearest neighbour search beyond the settings we discuss in our paper, including k-NN search, different distance metrics such as L1 distance, and structures that support insertions and deletions. Moreover, the frequency estimation setting could be easily adapted to cover membership data structures such as bloom filters.
For new data structure problems, such those mention in Section 3.2 could certainly the same architectures we have used, i.e. for the data-processing model, using a transformer for the batch-setting MLP for the streaming setting and using MLPs for the query-model. These architectures are relatively general and therefore require minimal adaptation. To be more specific, discuss how each of the problems we propose in Section 3.2 can be tackled with our framework. We focus on the batch setting as this is the more common one though similar insights apply to the streaming setting.
Graph data structures A cotnnutttion-qquvvvtttt architecture such as a transformer (with no positional encodings) or graph-neure network should be used to represent the graph the data- processing network. Note that it has already become commonplace to use transformers to represent graphs [56] [57. The output size of the model would be constrained to control the space complexity of the learned data structure. This is accomplished by either ignoring certain output tokens when less space then the input size is desired adding additional tokens when more space is required, as we do in our extra space experiments. The same MLP query architecture proposed should suffice.
37

---

Sparse matrices As matrices are simply grids of numbers, transformer would be appropriate as well with each token representing either row/column of the matrix single element. The output size of the model would be constrained to control the space-complexity of the learned data structure. The same MIP query architecture we proposed should suffice.
Learning statistical models In this case, the patt--prcessing network which operates over datasets should likely permutation-invariaat as there is no canonical ordering to an IID dataset so transformer would suffice. Again, the output size of the model would be constrained to control the size of the learned model (e.g. decision tree). The MLP query architecture would an appropriate choice for the query model.
While specialized architectures do exist for different domains (e.g. graph neural networks), trans- formers can certainly be used as starting point and are currently applied to many different types of inputs [58 59 56 There is also work that trying to build general architectures for structured inputs/outputs [60] 61. Advances in this direction could lead to more general architectures for data
structures as well
In summary, our framework is general in so far as it provides concrete way to set up the interaction between the data structure and the queries while respecting space and query complexity constraints. does not require prior knowledge of specific data structures or algorithms. That said, this should be viewed starting point-no an all-ennoopassing recipe-and and using problem-specific information can certainly improve performance.
E.2 Nearest Neighbor Permutation
A key NN constraint is that the prediction must be one of the original data points, which every NN method (even approximate ones) satisfies. Initially, we trained a transformer to output a re-ordered dataset directly, but training Instead, we trained the model to learn a permutation directly, avoiding the need to learn to reproduce the data while still preserving expressivity, since any structure (e.g. clustering hierarchical ordering) can be represented as a permutation. Our experiments, such as the FashionMNIST plot in Figure [8 confirm that this method naturally discovers clustering. Beyond NN, it can benefit any task requiring retention of the inputs.
slower because it had to recreate the original data as well reorder it.
38

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, the abstract and introduction clearly states the claims made and also references results that provide evidence for these claims, We also discuss the scope and limitations in the introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Yes, the limitations are addressed Section
5
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
39

---

Answer:
[NA]
Justification:
We have no theoretical results in the paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: In addition releasing the code provide architecture and hyperparameter details Section B.I All of our datasets are either generated with our code or are publicly available (e.g. MNIST, FashioMNIST, SIFT).
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
40

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We share code for the experiments and data generation.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips //ipss.cc/pill wttpcccpppppppppeeeeeeeeeeeeeeee hotaa//gggudeee/eeeeabiiiisssss /Cc////oddiiisssseeeeeee ssiiiiiisiyyyyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
This information is detailed in Sections B.I and
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We use bootstrapped confidence intervals for the majority of our experiments, we also show the variability of the results across seeds in Section B.2
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
41

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: In section B.I we detail the compute resources used and the number of optimization steps.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics atttps://neurre http:/p......p.llccccccccccsuddd
Answer: [Yes]
Justification:
We have reviewed the code of ethics and found violations.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: As this is a more fundamental paper focused on discovering algorithms, there is no immediate societal impact of the work performed.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
42

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Neither our data models pose risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All data is publicly available has been appropriately acknow (i.e. CAIDA dataset).
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
43

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Yes, we release the code and include documentation.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
44

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
45


Reviews:
[{"review_id": "VroNi6XZPH", "reviewer": "Reviewer_LhkM", "summary": "This paper introduced a method for learning data structures for nearest-neighbor search and count estimation. The data processing and query execution networks are learned in an end-to-end manner. The data is processed by a NanoGPT-based transformer, which returns a score for each data point that can be used to order the points. The model uses a differentiable sort function during training, allowing it to optimize the ordering of the data specifically for nearest-neighbor search.\nThe query execution network performs queries sequentially, where each model takes the query and the history of previous selections to output a one-hot vector that selects a data point. Given a budget of 𝑀 lookups, the query execution model aims to find the closest data point at each step. The network is trained using softmax-based soft lookups with added noise to encourage sparse selections. During inference, the softmax output is replaced with a hardmax operation to produce the final one-hot encoding for exact lookups.\nThe authors have discussed 2 use cases with this method:\n1. Near neighbor search: For 1 dimensional data, 2 dimensional data and n dimensional data (with different distributions)\n2. Frequency estimation", "questions": "1. How does the comparison looks like with SOTA NN search methods (graph based search like HNSW)? \n2. Does number of lookups translates well to the query latency for the proposed approach? (Is it a linear relationship?)", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper has introduced a novel idea of learning a data structure via transformers. The learning framework allows the data structure to adapt to the data distribution automatically.\n2. The idea is well motivated and analysis is provided on the behavior of data structure becoming similar and yet better than existing data structures like binary search, KD trees and LSH. \n3. The model handles the challenge of learning with discrete lookups by using softmax with noise during training and switching to hard softmax at the inference.\n\nWeaknesses:\n1. The paper needs more detailed experiments and analysis to validate the novelty of the idea. For example, the near neighbor search analysis will need more real world datasets (from ANN/BigANN benchmarks), RecallK@K measures, comparison of data structure sizes, construction (+learning) time and query latency (wall clock time and comparison on number of lookups).\n2. The paper doesn't give a clear picture on how the performance metrics vary with data dimensions.\n3. Other potential applications: The discussion in this section is vague. It lacks concrete examples or technical details for how this approach could be adapted to domains beyond nearest-neighbor search, making the section feel speculative and underdeveloped.", "was_revised": true, "final_justification": "Issues not resolved (but stated as not in the scope of the paper to claim)\n1. real world datasets (from ANN/BigANN benchmarks): \n2. RecallK@K measures: not resolved (accuracy is reported), \n3. Comparison with SOTA NN search methods (graph based search like HNSW)\n\nIssues resolved :\n1. Comparison of data structure sizes \n2. Construction (+learning) time and query latency (wall clock time and comparison on number of lookups).\n3. Performance metrics with data dimensions.\n4. Other potential applications discussion \n5. Does number of lookups translates well to the query latency for the proposed approach? (Is it a linear relationship?)", "rebuttal": "We thank the reviewer for your feedback and appreciate the time and effort you’ve dedicated to evaluating our work! Before responding to specific points, we’d like to clarify the high-level motivation behind our work.\n\nWe definitely understand the concerns around the efficiency and scalability of our methods and acknowledge that it’s necessary for future work to address this in order to build a system suitable for production workloads. Our current methods are not designed to scale to large datasets and we believe that more problem-specific inductive biases are crucial for this (see App D for a discussion).\n\nWe’d like to emphasize, however, that **our intention with this work was not to propose a method built for real-world deployment. Instead, we are interested in understanding to what extent neural networks can explore the space of data structures—even at small scale—from a classical algorithms perspective.**\n\nInitially, it was unclear (at least to us) whether end-to-end learning of data structures for tasks like nearest neighbors was even feasible, or how to even frame the problem. It raises many fundamental questions: How should data structures and queries be represented using neural networks, especially under combinatorial constraints like limited query budgets? Will training succeed at all, given that the data-processing and query networks are trained from scratch and rely entirely on each other? Even then, can such a system explore the space of classical and novel algorithms?\n\nTo our surprise, once we carefully set up training—with space and query constraints and a structured lookup interface—the model was able to discover meaningful algorithms! It recovered classic data structures like sorting, KD-trees, and LSH-like behavior, often adapting them to the underlying distribution. \n\nTo better understand what the model is learning, we place a strong emphasis on reverse-engineering its behavior and comparing it to classical algorithms. This perspective also motivates our focus on theoretical notions of efficiency such as query and space complexity—quantities we can directly enforce and optimize during training. In contrast, metrics like wall-clock time are often closely tied to hardware and implementation details and were not the central focus of our investigation.\n\nTo summarize, we view this as our core contribution: framing data structure discovery as an end-to-end learning problem and showing that it can explore a rich algorithmic space.  Finally, to our knowledge, this is one of the first attempts to approach data structure discovery in a general way. While our experiments focus on NN search, the design principles extend to other algorithmic settings. Of course, scaling to real-world workloads is an exciting next step, but our results show that even small-scale models can yield useful insights—potentially informing the design of scalable, hand-engineered systems\n\nWe hope this clarifies our intentions with this work. Nevertheless, we’ve measured the throughput of our models for the nearest neighbor experiments in several regimes.\n\n\n### Performance Metrics\nWe emphasize that we did not optimize for practical efficiency as we were primarily interested in the fundamental tradeoffs between space and query complexity, though these are certainly related to runtime in practice. We use the same architectures from our paper (see App B.1 for details). Benchmarking is conducted with a batch size of 512 and the results are averaged over 10 trials on a single GPU. \n\n**Query Processing Throughput**\n\nNumber of queries processed per second when using log(n) lookups per query. As expected, querying smaller datasets has higher throughput as the number of lookups, and thus neural net forward inferences, is smaller.\n\n| dim  | n=64   | n=128  | n=256  | n=1024 | n=2048 |\n|------|--------|--------|--------|--------|--------|\n| 32   | 212943 | 175559 | 146495 | 113960 | 103476 |\n| 64   | 211518 | 173807 | 149603 | 113249 | 105833 |\n| 128  | 208741 | 172275 | 151533 | 116607 | 104858 |\n\nBelow, we fix N=256 and measure how increasing the number of lookups (M) affects throughput. There’s a clear inverse correlation between M and throughput indicating that query complexity has implications for practical efficiency.\n\n| dim  | M=1     | M=2    | M=3    | M=4    | M=5    | M=6    | M=7    | M=8    |\n|------|---------|--------|--------|--------|--------|--------|--------|--------|\n| 32   | 1072476  | 517067 | 345946 | 261705 | 212202 | 176503 | 145927 | 134165 |\n| 64   | 1090290 | 513129 | 355654 | 265753 | 211991 | 174031 | 150917 | 132594 |\n| 128  | 1054366 | 519691 | 348347 | 260985 | 212237 | 173406 | 145191 | 133278 |\n\n**Data-preprocessing Throughput**\n\nThe table below shows the number of raw datasets processed into a data structure per second using quadratic and linear attention (which we explored as a faster alternative, see App. D.1).  Quadratic attention runs faster here due to flash-attention [7], which optimizes memory access. Notably, in smaller N regimes, linear attention can be less efficient than quadratic attention with hardware optimizations, despite its lower FLOP count [7].\n\n| Arch         | Dim | n=64  | n=128 | n=256 | n=1024 | n=2048 |\n|--------------|-----|-------|-------|-------|--------|--------|\n| **Linear**   | 32  | 17701 | 8772  | 4542  | 1170   | 585    |\n| Linear       | 64  | 17241 | 8844  | 4563  | 1177   | 590    |\n| Linear       | 128 | 17121 | 8614  | 4473  | 1152   | 577    |\n| **Quadratic**| 32  | 95229 | 49697 | 13418 | 5867   | 571    |\n| Quadratic    | 64  | 94815 | 46924 | 13285 | 5248   | 548    |\n| Quadratic    | 128 | 93294 | 48506 | 13095 | 5174   | 543    |\n\nWe also try fixing the input dimension to 128 and adjust the number of layers in the quadratic transformer.\n\n| # Layer | n=64   | n=128  | n=256  | n=1024 | n=2048 |\n|---------|--------|--------|--------|--------|--------|\n| 4       | 174433  | 174381 | 173501 | 31512  | 8449   |\n| 8       | 94493  | 95037  | 18849  | 6780   | 610    |\n| 12      | 65331  | 24301  | 8073   | 1103   | 220    |\n\n**Training Time**\n\nBelow is the training time for each of our experiments with N=100 sized datasets\n| 1D        | 2D        | 30D        | MNIST      | SIFT       | FashionMNIST |\n|-----------|-----------|------------|------------|------------|--------------|\n| ~3 Hours  | ~3 Hours  | ~10 Hours  | ~10 Hours  | ~2 Hours   | ~2 Hours     |\n\n\n>How does the comparison looks like with SOTA NN search methods (graph based search like HNSW)?\n\nAs discussed in the overall motivation for our work above, we do not include comparisons to SOTA NN methods such as HNSW as these are complex systems explicitly designed to maximize runtime efficiency in practice which is not the intention of our work.\n\n>Does number of lookups translates well to the query latency for the proposed approach? (Is it a linear relationship?)\n\nYes! The number of lookups is inversely proportional to query throughput (see tables provided above). Although the relationship is not exactly linear in our measurements this is likely due to some constant overhead as the relationship should asymptotically be linear. This is because each additional lookup requires an additional forward pass through a query model.\n\n\n### Other Potential Applications\n>Other potential applications: The discussion in this section is vague. It lacks concrete examples or technical details for how this approach could be adapted to domains beyond nearest-neighbor search, making the section feel speculative and underdeveloped.\n\nThanks for your feedback. Our goal in the “Other Potential Applications” section is simply to show that a wide range of problems—beyond nearest‑neighbor search—can be framed as “data‑structure discovery” and tackled with similar methods to those which we use.  We do not go into too much depth as a full treatment of each would indeed deserve its own paper. However, we kindly ask the reviewer to refer to appendix E where we have provided more detail on how the framework can be applied to these other problems including technical details such as architecture choices. As an example, we can describe how to adapt the framework to sparse matrix problems. Here, the goal is to efficiently answer queries about an NxN matrix (e.g., a matrix-vector product). The data network compresses the matrix into K tokens (with K < N); the query model retrieves M data chunks based on a query vector, and the predictor network computes the final response (e.g., learning to perform the product). The same architectures from NN search and frequency estimation can be used for this problem. Please let us know if you’d like additional clarification. \n\n\nWe hope this response answers your questions and addresses your concerns. Please let us know if you have any other questions!"}, {"review_id": "BTXSodoPTq", "reviewer": "Reviewer_aUtU", "summary": "This paper is primarily about learning data structures for fast nearest neighbor search.\n\nTwo models are learned:\n\n1. A transformer that takes as input a data set (consisting of N vectors, each d-dimensional) and outputs a single scalar for each data point. The points are then sorted by this scalar value.\n\nIn other words, the transformer selects an *ordering* of the points. The is the *data structure*.\n\n2. A collection of M multilayer perceptrons that take a query point and select possible candidates from the ordered collection of points. This generates a set of M candidates, and of these the nearest (in Euclidean distance) is returned.\n\nTo be precise, the jth MLP takes as input the query point as well as the results of the first j-1 MLPs; it then outputs a position (1 to N) in the ordering, identifying the data point at this position.\n\nVarious experiments are conducted:\n\n(a) One dimensional data that is uniformly distributed over [-1,1].\n\nThe authors take N=100 and M=7. They find that the data structure (approximately) sorts the points in increasing order. Meanwhile the query component does better than binary search: given a query point q, it can approximately predict its position in the ordering, and uses this in the search.\n\n(b) One dimensional data with a harder distribution.\n\nOnce again, the data structure ends up sorting the data, but this time the query component uses something like binary search.\n\n(c) Two dimensional data distributed uniformly over [-1,1]^2.\n\nThis time, the data structure produces an ordering of the points that is approximately locality-preserving. The authors find it similar to a k-d tree.\n\n(d) Two dimensional data from a harder distribution.\n\nAgain, the authors describe the data structure as akin to a k-d tree.\n\n(e) High-dimensional data, e.g. Fashion MNIST.\n\nHere the authors relate the learned data structures to locality-sensitive hashing.\n\nFinally, the authors discuss other applications of the methodology, including a detailed application to frequency estimation in data streams.", "questions": "1. In the 1-d case with the uniform distribution:\n\n(a) The paper says that the data structure is 99.5% accurate in perfectly ordering the data. Since N=100, this means that about half the time, the data is perfectly ordered, and half the time, there is one point out of place? Is that correct? How far out of place is that one point?\n\n(b) It would be nice to see in more detail what the query engine is doing. Does it map query q to position (q+1)/2 * N, and then do something like a binary search around that location? I think it would be helpful to see a figure in which, for a particular query (e.g. q = 0.7), the positions of the M probes are shown. This would give a better sense of the querying strategy. In general, the paper does a reasonable job of analyzing the data structures, but not so much the querying strategies.\n\n(c) In this case, although the data structure needs to be relearned for each new data set (from this distribution), it should be possible to keep the querying module fixed... is that right?\n\n2. In the 2-d case with the uniform distribution:\n\nThe data structure is some sort of space-filling curve? It would be interesting to see at least a little bit of this curve. The pictures in Fig 9 are great, but the color coding makes it hard to see detail (i.e. there aren't 100 different colors, so the exact ordering in unclear).\n\nFor the querying strategy, I really like Fig 10. In fact, Figs 9 and 10 are the most interesting things in the paper and so it is a shame they are in the appendix!", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "The good:\n\n1. The overall question (of learning data structure and query engine end-to-end) is fascinating.\n\n2. The choice of nearest neighbor is perfect, as this is a problem with a long history of clever-yet-simple data structures.\n\n3. The choice of data structure -- namely, a permutation of the data -- is sensible.\n\n4. The analysis of the data structures in the 1-d and 2-d setting is quite careful and therefore interesting. In particular, the pictures of the 2-d ordering (in the appendix) are enlightening.\n\nThe not-so-good:\n\n1. The 1-d and 2-d data structures and query engines should have been examined in even more detail (see questions later).\n\n2. The higher-dimensional examples have too little detail to be of interest.\n\n3. Likewise with the frequency estimation and other stuff; I think it would be have preferable to skip all of this and just spend more time on the nearest neighbor case.", "was_revised": true, "final_justification": "The authors have answered my questions very thoroughly. I maintain my original positive score. I find this paper to be creative, carefully crafted, and insightful. In particular:\n\n* The experimental design makes it possible to carefully study the (neural) data structure and query strategy, and the authors do a good job of this in the 1-d and 2-d settings.\n* The learned strategies are fascinating hodgepodges of many data structures that have been developed for NN search over the decades.\n* The strategies are a reminder that even in the simplest of cases (e.g. using binary search to identify an element in a sorted list), there are significant gains in efficiency possible from taking the data distribution into account (e.g. permitting a preliminary \"guess\" of where the item might be).", "rebuttal": "We thank the reviewer for your feedback and appreciate the time and effort you’ve dedicated to evaluating our work! We’re glad to hear you found the direction fascinating and enjoyed the visualizations of the data and query processing! Please find our responses to your questions below. Unfortunately, as we can’t upload a new draft or share any links we cannot provide figures for some of your questions. However, we provide our findings below and will definitely update the next draft to include the corresponding figures. \n\n>(a) The paper says that the data structure is 99.5% accurate in perfectly ordering the data. Since N=100, this means that about half the time, the data is perfectly ordered, and half the time, there is one point out of place? Is that correct? How far out of place is that one point?\n\nYes, your understanding is correct! We did some additional analysis and found that when the model does make a sorting error, the two points that need to be swapped are neighbors. \n\n>(b) It would be nice to see in more detail what the query engine is doing. Does it map query q to position (q+1)/2 * N, and then do something like a binary search around that location? I think it would be helpful to see a figure in which, for a particular query (e.g. q = 0.7), the positions of the M probes are shown. This would give a better sense of the querying strategy. In general, the paper does a reasonable job of analyzing the data structures, but not so much the querying strategies.\n\nIn Figure 2 (center) we show the query model’s behavior on the first lookup. For the 1D uniform distribution, we found that the model starts its search based on its prior (i.e. if q=0.7 and N=100 the model starts at position 70) and then does a local search in that neighborhood in a fashion similar to interpolation search (rather than binary search). \n\n>(c) In this case, although the data structure needs to be relearned for each new data set (from this distribution), it should be possible to keep the querying module fixed... is that right?\n\nOnce both networks are trained on datasets/queries from the relevant distributions their parameters no longer need to be updated. After training, given a new dataset (from the same distribution), we simply run a single forward pass through the data-processing network to construct the data structure. The data structure can then be reused multiple times to respond to queries, processed by the query networks. I hope this answers your question, please let us know if you’d like additional clarification!\n\n>The data structure is some sort of space-filling curve? It would be interesting to see at least a little bit of this curve. The pictures in Fig 9 are great, but the color coding makes it hard to see detail (i.e. there aren't 100 different colors, so the exact ordering in unclear).\n\nThat’s a great question and yes your intuition is correct! We plotted the data structure model’s decision boundary and we see that it roughly divides the space into evenly sized partitions.\n\nWe hope this response answers your questions and addresses your concerns. Please let us know if you have any other questions!"}, {"review_id": "WKrGbLHxe4", "reviewer": "Reviewer_LJYb", "summary": "This paper investigates the feasibility of learning data structures end-to-end using neural networks, with a primary focus on the nearest neighbor search problem in both low and high dimensions. The authors propose a novel framework consisting of two jointly trained components: a data-processing network that constructs a data structure from raw input data, and a query network that learns how to query it efficiently.\nThe study demonstrates that, neural networks can rediscover classical algorithms such as sorting followed by binary search, and in higher dimensions, structures resembling k-d trees or locality-sensitive hashing. The framework is extended to a frequency estimation task in data streams, showcasing its broader applicability. The paper highlights both the potential and challenges of this learning paradigm and provides empirical evidence that neural models can autonomously learn interpretable and effective data structures from scratch.", "questions": "1.  Please address W1 by discussing relevant prior work on end-to-end neural data structures learning, and explaining how the proposed approach differs from or improves upon these works. It would also be helpful to include experimental comparisons to highlight the novelty or advantages of the proposed method.\n\n2.  Please support or clarify W2 with additional experiments, preferably by providing efficiency-related metrics such as runtime, throughput, or similar comparisons.\n\n3.  For W3, it would be helpful if the authors could include empirical measurements of preprocessing and training time across different model sizes, as well as how these metrics scale with the size of the dataset.\n\n4.  For W4, although the authors mention that worst-case guarantees cannot be provided, it would still be valuable if they could offer any empirical observations or partial theoretical guarantees.", "limitations": "Yes", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths：\n\n1. The paper introduces an interested idea—treating data structure design as an end-to-end learning problem.\n\n2. The proposed framework is effectively applied to classic tasks such as nearest neighbor search and frequency estimation. It can outperform classical data structures under uniform distribution and demonstrate comparable performance under hard distribution.\n\n3.The authors provide thorough ablations that isolate the contributions of different components, accompanied by well-reasoned explanations.\n\n4.The paper is well written and clearly organized.\n\nWeaknesses：\n\n1. W1(Related Works Missing). The paper seems to miss some important prior work on end-to-end learning of data structures. As far as I know, [1] introduced the CLRS Algorithmic Reasoning Benchmark, which also trains graph neural networks in an end-to-end manner to perform dozens of classical algorithms such as sorting and searching. That work, along with follow-up studies like [2, 3], should have been mentioned. In addition, for the task of frequency estimation, there are also recent neural data structure approaches based on end-to-end learning, such as [4, 5, 6]. These cutting-edge and highly relevant works are not discussed or compared against in the paper, which feels like a significant oversight.\n\n2. W2(Efficiency Overheads). The approach often relies on large Transformer-based architectures, leading to significant pre-processing and inference costs. The quadratic complexity of standard attention limits applicability without further optimization.\n\n3. W3(Scalability Limitations). Most experiments are conducted on small datasets, raising concerns about the framework’s feasibility for large-scale or real-world applications. Scaling to practical-sized problems remains an open challenge.\n\n4. W4(Guarantee Absence). Despite the strong empirical results, the absence of worst-case guarantees is a notable limitation—especially important when comparing to well-understood classical algorithms.\n\n[1] The CLRS Algorithmic Reasoning Benchmark. ICML\n\n[2] Neural Algorithmic Reasoning with Causal Regularisation. ICML\n\n[3] Neural Algorithmic Reasoning Without Intermediate Supervision. NeurIPS\n\n[4] Meta-sketch: A Neural Data Structure for Estimating Item Frequencies of Data Streams. AAAI\n\n[5] Learning to Sketch: A Neural Approach to Item Frequency Estimation in Streaming Data. TPAMI\n\n[6] Lego Sketch: A Scalable Memory-augmented Neural Network for Sketching Data Streams. ICML\n\nIf these issues can be resolved, I will consider further improving my score.", "was_revised": true, "final_justification": "My primary concern with this paper is that it overlooks two important lines of related work (the CLRS benchmark and Neural Sketches) , which were not discussed or experimentally compared at all in the original version. This omission poses a significant challenge to both the completeness and the novelty of the work.\n\nWhile I accept that the CLRS benchmark need not be included in experiments due to substantial technical differences, it should still be discussed as representative work in the field. The omission of Neural Sketches is far more critical: despite acknowledging its close technical and conceptual similarity to this paper, the authors argue it should not be compared, citing a focus on nearest-neighbor search. However, this task is inherently supported by memory-augmented neural networks, a capability central to the Neural Sketches series of works. This omission undermines the paper’s completeness and scholarly positioning, and cannot be addressed within the scope of a camera-ready version.", "rebuttal": "Thank you very much for your detailed review and for providing useful feedback!\n## W1: Relation to Prior Work\nThanks for sharing these references! We comment on them below and will definitely include a more thorough discussion on these in our next draft.\n\n**CLRS**\n\nWhile works related to the CLRS benchmark [1, 2, 3] share a similar spirit to ours of studying neural networks (NNs) and classical algorithms there are several key differences:\n- These works train NNs to simulate specific algorithms (e.g. bellman-ford). We aren’t trying to learn any specific algorithm but rather want the learned algorithms to satisfy certain conditions (i.e. space and query complexity constraints). \n- Our setup uses significantly less supervision than these works which usually provide the model with intermediate outputs of the algorithm they’re trying to learn as supervision. While [3] does not use this supervision they do add additional inductive biases such as a contrastive learning objective to make the problem more tractable. \n- Our setup is more challenging as the model needs to learn two algorithms in tandem (data processing and query algorithms) from scratch. \n- These works tend to use specialized graph NNs whereas we instead introduce minimal problem-specific biases, aiming to test whether general-purpose NNs can learn data structures.\n\n**Neural Sketching**\n\nWe were not aware of this recent work on neural sketching algorithms and thank the reviewer for pointing it out! Their approach shares several similarities with our frequency estimation experiments—both can be viewed as memory-augmented neural networks trained to read from and write to differentiable memory [8]. While we did not run direct comparisons, their architectures and training setups suggest performance comparable to, and possibly better than, ours, as they were tailored specifically for frequency estimation. In contrast, our experiments aimed to show the generality of our framework, with our main focus on nearest-neighbor search.\n\nWe also highlight a difference in motivation. Their work targets scalable, practically efficient neural sketches, whereas ours emphasizes theoretical notions of efficiency, such as query complexity. For instance, while their work encourages sparse read/writes we carefully enforce one-hot query sparsity in order to control the number of lookups. This design choice lets us draw direct comparisons to classical sketches (e.g. Count‑Min Sketch) in terms of query complexity, rather than solely throughput or runtime.\n\nFinally, while their work focuses on frequency estimation, our work is more concerned with identifying the broader principles behind data structure discovery than trying to learn practical data structures for a single problem. For instance, it’s not clear how to translate their framework to also handle NN search or some of the other problems we discuss. We will include these references in the next draft, discuss their connections to our work, and either provide a comparison or acknowledge that their method likely performs similarly or better for this task.\n## W2/W3: Efficiency & Scale\n\nWe definitely understand the concerns around the efficiency and scalability of our methods and acknowledge that it’s necessary for future work to address this in order to build a system suitable for production workloads. Our current methods are not designed to scale to large datasets and we believe that more problem-specific inductive biases are crucial for this (see App D for a discussion).\n\nWe’d like to emphasize, however, that **our intention with this work was not to propose a method built for real-world deployment. Instead, we are interested in understanding to what extent neural networks can explore the space of data structures—even at small scale—from a classical algorithms perspective.**\n\nInitially, it was unclear (at least to us) whether end-to-end learning of data structures for tasks like nearest neighbors was even feasible, or how to even frame the problem. It raises many fundamental questions: How should data structures and queries be represented using neural networks, especially under combinatorial constraints like limited query budgets? Will training succeed at all, given that the data-processing and query networks are trained from scratch and rely entirely on each other? Even then, can such a system explore the space of classical and novel algorithms?\n\nTo our surprise, once we carefully set up training—with space and query constraints and a structured lookup interface—the model was able to discover meaningful algorithms! It recovered classic data structures like sorting, KD-trees, and LSH-like behavior, often adapting them to the underlying distribution. \n\nTo better understand what the model is learning, we place a strong emphasis on reverse-engineering its behavior and comparing it to classical algorithms. This perspective also motivates our focus on theoretical notions of efficiency such as query and space complexity—quantities we can directly enforce and optimize during training. In contrast, metrics like wall-clock time are often closely tied to hardware and implementation details and were not the central focus of our investigation.\n\nTo summarize, we view this as our core contribution: framing data structure discovery as an end-to-end learning problem and showing that it can explore a rich algorithmic space.  Finally, to our knowledge, this is one of the first attempts to approach data structure discovery in a general way. While our experiments focus on NN search, the design principles extend to other algorithmic settings. Of course, scaling to real-world workloads is an exciting next step, but our results show that even small-scale models can yield useful insights—potentially informing the design of scalable, hand-engineered systems\n\nWe hope this clarifies our intentions with this work. Nevertheless, we’ve measured the throughput of our models for the nearest neighbor experiments in several regimes.\n\n### Performance Metrics\nWe emphasize that we did not optimize for practical efficiency as we were primarily interested in the fundamental tradeoffs between space and query complexity, though these are certainly related to runtime in practice. We use the same architectures from our paper (see App B.1 for details). Benchmarking is conducted with a batch size of 512 and the results are averaged over 10 trials on a single GPU. \n\n**Querying Throughput**\n\nNumber of queries processed per second when using log(n) lookups per query. As expected, querying smaller datasets has higher throughput as the # of lookups, and thus neural net inferences, is smaller.\n\n| dim  | n=64   | n=128  | n=256  | n=1024 | n=2048 |\n|------|--------|--------|--------|--------|--------|\n| 32   | 212943 | 175559 | 146495 | 113960 | 103476 |\n| 64   | 211518 | 173807 | 149603 | 113249 | 105833 |\n| 128  | 208741 | 172275 | 151533 | 116607 | 104858 |\n\n**Data-preprocessing Throughput**\n\nThe table below shows # datasets processed into a data structure per second using quadratic and linear attention (which we explored as a faster alternative in App. D.1). Quadratic attention runs faster here due to flash-attention [7], which optimizes memory access. Notably, in smaller N regimes, linear attention can be less efficient than quadratic attention with hardware optimizations, despite its lower FLOP count [7].\n\n| Arch         | Dim | n=64  | n=128 | n=256 | n=1024 | n=2048 |\n|--------------|-----|-------|-------|-------|--------|--------|\n| **Linear**   | 32  | 17701 | 8772  | 4542  | 1170   | 585    |\n| Linear       | 64  | 17241 | 8844  | 4563  | 1177   | 590    |\n| Linear       | 128 | 17121 | 8614  | 4473  | 1152   | 577    |\n| **Quadratic**| 32  | 95229 | 49697 | 13418 | 5867   | 571    |\n| Quadratic    | 64  | 94815 | 46924 | 13285 | 5248   | 548    |\n| Quadratic    | 128 | 93294 | 48506 | 13095 | 5174   | 543    |\n\nWe also try fixing the data dimension to 128 and adjust the # of layers in the quadratic transformer.\n\n| # Layer | n=64   | n=128  | n=256  | n=1024 | n=2048 |\n|---------|--------|--------|--------|--------|--------|\n| 4       | 174433  | 174381 | 173501 | 31512  | 8449   |\n| 8       | 94493  | 95037  | 18849  | 6780   | 610    |\n| 12      | 65331  | 24301  | 8073   | 1103   | 220    |\n\n**Train Time**\nBelow is the train time for each of our experiments with N=100 sized datasets\n| 1D        | 2D        | 30D        | MNIST      | SIFT       | FashionMNIST |\n|-----------|-----------|------------|------------|------------|--------------|\n| ~3 Hours  | ~3 Hours  | ~10 Hours  | ~10 Hours  | ~2 Hours   | ~2 Hours     |\n\n## W4: Guarantees\nThanks for raising this point! While our learned data structures lack formal worst-case guarantees, their empirical performance is encouraging and can provide some insight:\n- We find that in many regimes, the models learn algorithms resembling classical ones—e.g., binary search, interpolation search, KD trees, and LSH—in both structure and querying, indicating they’re approximating the same performance guarantees those algorithms enjoy. \n- Beyond 1D search, few methods offer worst-case guarantees better than O(n); for example, both LSH and KD trees have O(n) worst-case complexity. Our models can trivially match this by exhaustively searching, but more meaningfully, we show that their average-case performance is often comparable or better.\n- One could use our setup to derive insights that could then be translated into a learning-augmented algorithm that does have guarantees (see App C.3. for an example with frequency estimation), or combine it with an existing approach, for instance use a learned model for a quick search and fall back to a method with guarantees if nothing is found, thereby getting the best of both worlds.\n\nWe hope this response addresses your concerns. Please let us know if you have other questions!\n\n[7]  https://arxiv.org/abs/2205.14135\n\n[8] https://www.nature.com/articles/nature20101"}, {"review_id": "J0wNuNAFSF", "reviewer": "Reviewer_ffsc", "summary": "The paper proposes an end-to-end framework for data structure training for NN search tasks and beyond. The method comprises a differential sorting network to form the data structure and a query network to conduct the query. This framework works great in different NN search tasks with different data dimensions and distributions. The paper also explores its application on tasks beyond the NN search like frequency estimation.", "questions": "See weakness above.", "limitations": "yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths: \n1. The paper writing is fluent and easy to read, the motivation is easy to follow and the figures and tables are also clear.\n2. The method design is simple and effective, forming the data structure on the fly when training on query tasks.\n3. The experiment findings are very interesting, which forms different structures in different data dimensions and the structures align with some data structure algorithms. These findings also show the effectiveness of the method.\n4. The method can be used in tasks beyond NN search, like frequency estimation, and also show promising results.\n5. The experiments are very sufficient to prove the claims of the paper.\n\nWeaknesses:\n1. Could you describe how to adapt your model in frequency estimation in detail? I still don't understand the pipeline of this application.\n2. In the current age, the dimensions of 30 and 100 are still a little bit small. Have you tried your method on more dimensions?\n3. If the data processing model can try to struct the data not only in a 1D sorting, but in a 2D map or higher dimension, the learnt data structure may have more interesting properties and have more further applications.", "was_revised": true, "final_justification": "This is a solid and interesting work, and the authors' responses have addressed most of my concerns. Thus, I maintain my accept score.", "rebuttal": "We thank the reviewer for your feedback and appreciate the time and effort you’ve dedicated to evaluating our work!\n\n>Could you describe how to adapt your model in frequency estimation in detail? I still don't understand the pipeline of this application.\n\nUnfortunately, due to space constraints we were unable to provide more detail in the main body. However, in App C.2 we provide a more detailed description of how we adapt the model for frequency estimation. In Appendix E, we discuss the broader applicability of the framework and how it can be adapted to several other problems. Please let us know if there are any specific details that are still unclear and we’d be happy to respond! \n\n>In the current age, the dimensions of 30 and 100 are still a little bit small. Have you tried your method on more dimensions?\n\nYes, we include experiments with a 3-digit version of MNIST which has dimensions of 28*84 in Section 2.4. We find that the model learns a low-dimensional representation of the data, processes the data into a data structure and queries the data structure all in-tandem end-to-end. \n\n>If the data processing model can try to struct the data not only in a 1D sorting, but in a 2D map or higher dimension, the learnt data structure may have more interesting properties and have more further applications.\n\nWe definitely agree! As both models are being trained end-to-end the data-structure model effectively learns a useful space partitioning in high-dimensions that can be efficiently queries. For instance, in Figure 18 we show the model’s learned clustering of FashionMNIST data (100 Dimensions). \n\nWe hope this response answers your questions and addresses your concerns. Please let us know if you have any other questions!"}]

Author rebuttal:

Meta-review: This paper proposes a method to learn data structures end-to-end with neural networks, focusing on nearest-neighbor (NN) search but also including frequency estimation. It has received mostly positive reviews.

The paper was found clearly organized, well written, fluent, easy to read, with clear figures and tables. The problem was found challenging, fascinating, with high impact. The choice of NN search was found perfect. The method was found well motivated, novel, interesting, creative, sensible, simple, effective, technically solid. The analysis was found interesting, insightful. The experiments were found sufficient, with thorough ablations. The performance was found competitive with classical data structures. The findings were found very interesting

Concerns included missing related work, relying on large architectures incurring significant cost, scalability, insufficiency of large datasets, metrics and ablations, data dimensions being small, the section on other applications being vague and underdeveloped, lack of detail, the frequency estimation part unclear and recommended to skip, absence of worst-case guarantees

The authors provided detailed feedback in their rebuttal, accompanied with new experiments where needed. This feedback addressed most of reviewers' concerns and substantial discussion followed. The authors committed to updating the paper with the new material they prepared during rebuttal and discussion.

Concerns that remain after discussion include insufficiency of large datasets, lack of RecallK@K measures, comparison with SOTA NN search methods (Reviewer LhkM) and missing related work CLRS and Neural Sketches (Reviewer LJYb, who remained negative).

On the former, the authors argued that they do not to propose a method built for real-world deployment but rather to investigate to what extent neural networks can explore the space of data structures. The reviewer understood the motivation behind the work and remained positive.

On the latter, very long discussion followed between the authors and the reviewer, as well as the reviewer and myself. The relation to CLRS appears to have been addressed. However, the reviewer found that the relation to Neural Sketches, a recent work, is critical.

The authors accepted that there are similarities between their work and Neural Sketches, but also differences, such that the novelty of this work is not challenged. Their work focuses mainly on NN search while Neural Sketches are specifically designed for frequency estimation. They have committed to discuss Neural Sketches (as well as CRLS) and to "acknowledge that they [Neural Sketches] are a more scalable solution for frequency estimation". The reviewer has insisted that "the authors have overlooked the broader field of memory-augmented neural networks" and suggested that the author adapt Neural Sketches to NN search, which however does not appear to be straightforward or within the scope of this work.

Given the very positive position of the other reviewers and the investigative nature of this work on a challenging problem that has not been addressed before, I find that this work has merit for publication and I recommend that it is accepted. However, given the discussions, I recommend that the authors move "other applications" discussions (section 3.2) to the supplementary and then either add detailed design and comparison with Neural Sketches if possible in the frequency estimation part (section 3.1) or entirely move the frequency estimation part to the supplementary as well, leaving space more space for the new material and strengthening the NN part.
Decision: Accept (poster)