# Discovering Data Structures: Nearest Neighbor Search and Beyond

arXiv: https://arxiv.org/abs/2411.03253 · HTML: https://mlsanity.com/paper/2411.03253/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 36 of its 62 references that are in the map.
- First posted on arXiv 5 Nov 2024.

## Builds on (most cited first, 30 of 36)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [Point Transformer](https://mlsanity.com/paper/2011.00931/) (2020)
- [Perceiver: General Perception with Iterative Attention](https://mlsanity.com/paper/2103.03206/) (2021)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [A Generalization of Transformer Networks to Graphs](https://mlsanity.com/paper/2012.09699/) (2020)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://mlsanity.com/paper/2107.14795/) (2021)
- [Neural Combinatorial Optimization with Reinforcement Learning](https://mlsanity.com/paper/1611.09940/) (2016)
- [Learning a SAT Solver from Single-Bit Supervision](https://mlsanity.com/paper/1802.03685/) (2018)
- [The Case for Learned Index Structures](https://mlsanity.com/paper/1712.01208/) (2017)
- [Algorithms with Predictions](https://mlsanity.com/paper/2006.09123/) (2020)
- [ANN-Benchmarks: A Benchmarking Tool for Approximate Nearest Neighbor Algorithms](https://mlsanity.com/paper/1807.05614/) (2018)
- [Neural Execution of Graph Algorithms](https://mlsanity.com/paper/1910.10593/) (2019)
- [Differentiable Ranks and Sorting using Optimal Transport](https://mlsanity.com/paper/1905.11885/) (2019)
- [Stochastic Optimization of Sorting Networks via Continuous Relaxations](https://mlsanity.com/paper/1903.08850/) (2019)
- [The CLRS Algorithmic Reasoning Benchmark](https://mlsanity.com/paper/2205.15659/) (2022)
- [Approximate Nearest Neighbor Search in High Dimensions](https://mlsanity.com/paper/1806.09823/) (2018)
- [ALEX: An Updatable Adaptive Learned Index](https://mlsanity.com/paper/1905.08898/) (2019)
- [Learning to Hash for Indexing Big Data - A Survey](https://mlsanity.com/paper/1509.05472/) (2015)
- [Learning Augmented Binary Search Trees](https://mlsanity.com/paper/2206.12110/) (2022)
- [Neural Algorithmic Reasoning with Causal Regularisation](https://mlsanity.com/paper/2302.10258/) (2023)
- [Properly learning decision trees in almost polynomial time](https://mlsanity.com/paper/2109.00637/) (2021)
- [Neural Algorithmic Reasoning Without Intermediate Supervision](https://mlsanity.com/paper/2306.13411/) (2023)

## Abstract (arXiv)

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