Discovering Data Structures: Nearest Neighbor Search and Beyond

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

We excluded learning-to-hash methods as regular LSH is already a strong baseline in the 30D setup we considered and hashing-based methods are not designed to handle our MNIST setting. Specifically, in our MNIST setup the metric is not given but implicitly provided via nearest neighbor supervision. This means our model needs to learn the metric space, which learning-to-hash methods are not designed for. However, in response to your question and to highlight our models strengths in this regime, we compare with an optimal hash function - 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. To execute a NN search, a query is mapped to its corresponding bucket and the element closest to this number (defined over this 1D label space) is chosen as the NN. When N=50 and M=7 we find that the optimal K is 8. To be clear, this is an unfair comparison as our model has to learn features that capture the underlying metric space but in this case it's provided directly. However, this optimal hashing setting should still underperform our model in this regime as it does not use an adaptive query algorithm which is necessary for achieving the O(log N) performance our model recovers. See Figure 10 in our updated revision for the comparison between our model and this hashing baseline (LSH).

This raises an important point we’d like to emphasize about our model - by directly learning features, a data structure and a query algorithm end-to-end it can better exploit the low-dimensional structure of this problem. Specifically, learning-to-hash methods simply cluster the points by their hash code and then search every element within the corresponding hash bucket. Our method, however, can learn an adaptive query algorithm (in this case something like 1D interpolation search over the learned features), thereby identifying the NN with fewer lookups than a hashing-based method.

We do, however, include comparisons to learning-to-hash methods for the additional high-dimensional SIFT/FashionMNIST experiments we ran as these are fairer comparisons (See Appendix G.2 of the paper).

The application of the framework to nearest neighbor search and frequency estimation seems somewhat disconnected. Could the authors elaborate on the principles that unify these applications within the framework? Could the authors provide a set of design principles or guidelines for adapting the framework to other tasks? This would make it clearer how researchers can extend this work. (W5)

Among other reasons, we chose to apply our framework to both of these settings to highlight the versatility of the framework and demonstrate how different data structure problems can be modeled by the same framework. In both settings, raw data first needs to be added to a data structure and then the data structure needs to be efficiently accessed to respond to some query. In our nearest neighbor setup, all the raw data is ingested together at one time (batch setting). However, in the frequency estimation setting, a single raw data point arrives at a time (streaming setting). Our framework is general and can handle both the batch and streaming setting, though different model architectures may be more appropriate for a given setting. For instance, in the streaming setting we use an MLP as the data-processing network as there is only a single point to process.  Both problems then require efficiently querying the data structure which we enforce by restricting the number of lookups the query models can make into the data structure. 

The framework is general in that it is designed to handle problems that require efficiently responding to queries about some collection of data. The queries access a data structure and efficiency is controlled by trading off the query complexity (number of lookups into the data structure) and the space complexity (size of the data structure). If other researchers had a candidate problem that could be framed like so (In section 3.2  we provide several examples), we believe our framework could be applicable. Specifically, the problem can be modeled with both a data-processing network and query network and could use similar principles we develop to control the efficiency of the data structure (e.g. tokenizing and restricting its size to control space complexity and using similar sparsity techniques on lookup vectors to enforce query complexity). The choice of architecture for each network would depend on the specific setting. For instance, for graph data-structures it may be more efficient to use a graph neural network for the data-processing network.

We hope this response answers your questions and addresses your concerns. Please let us know if you have any other questions!

We thank the reviewer for your detailed feedback and appreciate the time and effort you’ve dedicated to evaluating our work! 

We’ve addressed your specific questions below, however, we also kindly ask you to read our general response to all reviewers where we clarify the motivation and contributions of our work, discussing how it can provide useful insights for data structure design and avenues for scaling.

Considering the NN search problem has a trivial solution with O(n)complexity, this constraint raises concerns about the framework's practicality in real-world applications where much larger datasets are typical.

---

We’d like to clarify what may be a misunderstanding. For a given dataset, the data-processing network only needs to be executed once with an O(n^2) cost (this could be made even cheaper as we discuss in our next response), i.e. there is a fixed construction cost. Thereafter, the query models are restricted to M lookups where M < N and so each query to the data structure costs O(M).

Could the author discuss the justification for why a transformer with O(n2) complexity is necessary, and why using such a network is still feasible and beneficial for the NN search task? (W1)

We chose to use a transformer with O(n^2) complexity in order to keep the model relatively general. As you’ve noted, this may not always be necessary. For instance, high-dimensional hashing-based methods could be learned with a less complex model. However, data structures like k-d trees have a higher construction complexity O(dnlogn). Our framework does not require using any specific model so a less (or more) complex model can trivially be substituted in for the transformer. As we mention in our general response on scaling, one way to enable scaling to larger instances would be to lower the complexity of the data-processing model (e.g. using a linear transformer). 

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

Yes, we include additional experiments run on the SIFT (128 dim) and Fashion-MNIST (100 dim) datasets and compare them to learning-to-hash methods. Our model performs competitively with these learning-to-hash baselines, only slightly underperforming them. We didn’t do any hyperparameter tuning in order to respond quickly to reviewers but we’ll obtain final results (with tuning) for the camera-ready version. We also show an example of how the data-processing model learns to transform the Fashion-MNIST data by organizing the images by class. 

In addition to the HighD NN experiments, we also include an additional frequency estimation experiment where we show how insights we learned from our model suggested improvements to the CountMinSketch algorithm which lead to improvements on a large-scale real world IP dataset.

Please see Appendix G of the updated revision for more details on both the NN and frequency estimation experiments.

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

• (1) The paper uses NanoGPT, a decoder-only transformer that typically requires discrete tokens for input and output. However, the paper doesn’t clearly explain how continuous data inputs are tokenized or how outputs are transformed into the scalar ranking oi values. More detail on the embedding and decoding process for continuous data is necessary to understand the data flow and model structure.

• We apologize for the confusion. There is no specific reason as to why we chose NanoGPT other than that there is an open source implementation available that uses flash attention which allows for faster training.  Any other transformer model should work. Instead of using a discrete tokenizer, we have input and output projection layers (Linear layer) that project the input to the hidden dimension (64) and from the hidden dimension to the scalar values.

• (2) In experiments involving CNNs, essential details about the CNN architecture, such as the number of layers, kernel sizes, or feature extraction strategies, are missing. Including these specifications would make it easier to interpret the model's capabilities and constraints, especially regarding performance in high-dimensional setting**

• We use a simple CNN model with the following architecture:

<!---->

Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1)
ReLU()
MaxPool2d(kernel_size=2, stride=2, padding=0)
Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)
ReLU()
MaxPool2d(kernel_size=2, stride=2, padding=0)
Linear(in_features=64*7*21, out_features=128)
ReLU()
Linear(in_features=128, out_features=1)

We plan to fully open source the code before publication.

Thank you for your response! We address your remaining concerns below.

However, I remain cautious about the generalizability of the framework (W5), as the dependency on domain-specific design choices limits its broader applicability.

While we have already discussed how higher-level elements of our framework are applicable to new problems beyond those we explore, we’d like to clarify what would be required in terms of model architectures as we believe this may have been unclear from our prior responses.

For new data structure problems, such as those we mention in section 3.2, one could certainly use the same architectures we have used, i.e. for the data-processing model, using a transformer for batch-setting/MLP for streaming setting and using MLPs for the query-model. These architectures are relatively general and therefore require minimal adaptation. To be more specific, we discuss how each of the problems we propose in Section 3.2 can be tackled with our framework and minimal problem-specific knowledge (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 permutation-equivariant architecture such as a transformer (with no positional encodings) or a graph-neural network should be used to represent the graph in the data-processing network. Note that it has already become commonplace to use transformers to represent graphs [5, 6]. 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 or adding additional tokens when more space is required, as we do in our extra space experiments. The same MLP query architecture we proposed should suffice.

Sparse matrices: As matrices are simply grids of numbers, transformer would be appropriate as well with each token representing either a row/column of the matrix or a single element. The output size of the model would be constrained to control the space-complexity of the learned data structure. The same MLP query architecture we proposed should suffice.

Learning statistical models: In this case, the data-processing network which operates over datasets should likely be permutation-invariant as there is no canonical ordering to an IID dataset so a 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 be an appropriate choice for the query model.

In summary, we believe minimal problem-specific knowledge is required to adapt our framework to new types of data structure problems. While specialized architectures do exist for different domains (e.g. graph neural networks), transformers can certainly be used as a starting point and are currently applied to many different types of inputs [3, 4, 5]. There is also work [1, 2] that is trying to build general architectures for structured inputs/outputs. Advances in this direction could lead to more general architectures for data structures as well.

We will definitely include this more detailed discussion in our next revision.

Additionally, the lack of discussions of efficiency and scalability issues leaves critical questions unanswered for the applications requiring large-scale data.

In our original submission we included a section on scaling and limitations (Appendix F), however, we could certainly update this section to reflect our discussion with a more detailed analysis addressing avenues to improve efficiency as well (e.g. Linear Transformers). We will also move this section to the main paper to increase transparency. We will make these changes after the rebuttal period as currently we cannot upload a new revision.

We hope this addresses your concerns. Please let us know if you have any other questions.

[1] Jaegle, A., Gimeno, F., Brock, A., Vinyals, O., Zisserman, A., & Carreira, J. (2021, July). Perceiver: General perception with iterative attention. In International conference on machine learning (pp. 4651-4664). PMLR.

[2] Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula et al. "Perceiver io: A general architecture for structured inputs & outputs." arXiv preprint arXiv:2107.14795 (2021).

[3] Zhao, H., et al (2021). Point transformer. In Proceedings of the IEEE/CVF international conference on computer vision (pp. 16259-16268).

[4] Dosovitskiy, A. (2020). An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929.

[5] Dwivedi, V. P., & Bresson, X. (2020). A generalization of transformer networks to graphs. arXiv preprint arXiv:2012.09699.

[6] Ying, C. et al (2021). Do transformers really perform badly for graph representation?. Advances in neural information processing systems, 34, 28877-28888.

We thank the reviewer for your feedback and appreciate the time and effort you’ve dedicated to evaluating our work!

We kindly ask the reviewer to read the general response we posted as they address your comments regarding the scalability of learning data structures end-to-end.

Still requires a moderate amount of inductive bias: for example, the frequency estimation architecture is very different than the nearest neighbors one, and does not seem obvious to me. It seems that generalizing this approach to new data structure problems isn't entirely trivial and requires at least some trial-and-error, as well as expertise in how classical approaches to the data structure may work.

We acknowledge that the data-processing architectures are different for the NN and frequency estimation problems (the query architectures remain the same), and that other problems may benefit from problem-specific architectures however we still believe that the general framework we’re proposing can provide a scaffold for new data structure problems. Specifically, the problem can be modeled with both a data-processing network and query network and could use similar principles to control the efficiency of the data structure (e.g. tokenizing and restricting its size to control space complexity and using similar sparsity techniques on lookup vectors to enforce query complexity).

Do you have code to share to replicate your results?

We will definitely fully open-source the code before publication.

We hope this response answers your questions, please let us know if you have any other questions!

We thank the reviewer for your feedback and appreciate the time and effort you’ve dedicated to evaluating our work!

We kindly ask the reviewer to review the general response we posted as we believe there may be a misunderstanding around the motivation and contribution of our work. Our intentions with this work were neither to scale to very large datasets nor to beat existing SOTA approaches in these regimes (though we discuss how to further scale up the model in our response), but rather to understand what neural networks trying to learn data structures end-to-end are capable of doing, as a step towards using them for gaining insight into data structure design.

We hope this response addresses your concerns. Please let us know if you have any other questions!

Thank you for your response. As a reminder, we are not claiming that our specific method should be used as a direct substitute in existing data structure pipelines. Rather, we show that our framework can be used to gain insight into data structure design, and that these insights can be scaled and applied in realistic settings (e.g. our augmented CountMinSketch algorithm in Appendix G.1).

We’d like to re-emphasize that this is the first work that explores E2E learning of data structures (to the best of our knowledge) and so it seems premature to claim that they cannot be useful for advancing our understanding of data structures. In fact, augmenting data structures with predictions from neural networks has already been shown to be effective [1, 2, 3]. We’ve also discussed multiple avenues for future work to reduce the computational cost of learning data structures E2E.

[1] Tim Kraska, Alex Beutel, Ed H Chi, Jeffrey Dean, and Neoklis Polyzotis. The case for learned index structures. In Proceedings of the 2018 international conference on management of data, pp. 489–504, 2018.

[2] Honghao Lin, Tian Luo, and David Woodruff. Learning augmented binary search trees. In International Conference on Machine Learning, pp. 13431–13440. PMLR, 2022a.

[3] Yihe Dong, Piotr Indyk, Ilya Razenshteyn, and Tal Wagner. Learning space partitions for nearest neighbor search. arXiv preprint arXiv:1901.08544, 2019.

We thank the reviewer for your feedback and appreciate the time and effort you’ve dedicated to evaluating our work!

Q1. Is it possible to verify the effectiveness of the method with a large data that approaches the real applications?

We kindly ask the reviewer to read the general response we posted as they address your comments regarding the scalability and cost of learning data structures as well as the practical usefulness of our results. We also refer the reviewer to Appendix G of our updated revision where we include additional experiments showing how our model can provide useful, scalable insights and be applied to more realistic data.

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

Sorted lists, k-d trees, and LSH data structures are among the most widely used and fundamental data structures and so it is significant that our model can discover useful data structures. Moreover, beyond 1-dimension, data structure design is particularly challenging as there is no clear notion of ordering so recovering any useful data structure is impressive.

We also emphasize that it’s not trivial that any useful data structure can be recovered given that both the query and data structure model are being trained together from scratch - one can imagine that there exist many local minima where both models cannot make progress. For instance, in 1D, it is quite surprising that the query model can learn an algorithm such as binary search while also learning the correct data structure (a sorted list) to operate over. Moreover, both models are searching over a discrete and combinatorial space which poses even further optimization challenges.

We hope this response answers your questions and addresses your concerns. Please let us know if you have any other questions!

2. We show learning data structures end-to-end can provide useful insights into data structure design

As discussed above, the learned data structures can serve as a source of insight for data structure design. There are several points in this vein that are worth highlighting:

The learned data structures and insights do scale: While we don’t demonstrate the scalability of the model itself, the learned data structures and algorithms do in fact scale (e.g. LSH, KD Trees, Binary Search, etc). This suggests that one way one could use our framework is by first learning a data structure, followed by reverse-engineering it and applying it to larger settings.

Even when it is not possible to reverse-engineer the complete data structure, one can still discover insights to improve existing data structures. For instance, in the frequency estimation problem, we observed that the model performs well by using update deltas smaller than 1. We used this insight to create a modified version of CountMinSketch with smaller update deltas and applied it to a much larger and real world dataset and show that it outperforms traditional CountMinSketch by a factor of two in some regimes. (See Appendix G.1 of the updated revision for more details)

Distributional awareness: We show that distribution-dependent data structures and query algorithms emerge even at small scale. This suggests that one could use our framework to determine a lower-bound on performance for a distribution-dependent data structure.

Understanding the landscape of data structures: There are many fundamental open questions around understanding the landscape of possible data structures. For example, in low-dimensional nearest neighbor search, methods like KD-trees use O(n) space and guarantee O(log n) query complexity, where n is the number of input points. However, their query time scales exponentially with dimension. Conversely, methods like locality-sensitive hashing (LSH) for approximate nearest neighbor search in high dimensions use a reasonable amount of space and guarantee $O(n^{(1 - \delta)})$ query complexity, where \delta reflects the correlation between the query point and its nearest neighbor. While LSH avoids the exponential blow-up with dimension, the presence of $\delta$ in the exponent means performance degrades quickly as the correlation diminishes, unlike KD-trees. Are there any data structures in between that can achieve the best of both worlds? Our framework for end-to-end learning of data structures can possibly be used as a tool to make progress on such fundamental questions.

At the very least, one might hope that when high-dimensional data lies on low-dimensional manifolds (a common scenario), it would be possible to learn the low-dimensional structure and apply KD-tree-like methods to achieve logarithmic query complexity. We demonstrate that this is possible using our framework in the MNIST experiment, where our model identifies the relevant low-dimensional 1d feature in MNIST images (i.e. the relative ordering of numbers), sorts the data according to this feature, and applies binary search to achieve logarithmic query complexity.

3. Our proposed framework can be useful for modeling other data structure problems beyond NN/frequency estimation.

We believe that the general framework we’re proposing can provide a useful scaffold for other data structure problems beyond those that we examine in our work (in section 3.2 we provide several examples). Specifically, these types of problems can be modeled with both a data-processing network and query network and could use similar principles to control the efficiency of the data structure (e.g. tokenizing and restricting its size to control space complexity and using similar sparsity techniques on lookup vectors to enforce query complexity).