🤖 machine learning

Discovering Data Structures: Nearest Neighbor Search and Beyond

This paper proposes a general end-to-end learning framework that automatically discovers optimal data structures and query algorithms from scratch without initialization, successfully replicating known solutions like binary search, k-d trees, and locality-sensitive hashing for nearest neighbor search while also adapting to frequency estimation in data streams.

Original authors: Omar Salemohamed, Laurent Charlin, Shivam Garg, Vatsal Sharan, Gregory Valiant

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Omar Salemohamed, Laurent Charlin, Shivam Garg, Vatsal Sharan, Gregory Valiant

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive, messy library of books. Traditionally, librarians (computer scientists) spend years designing specific rules and filing systems (data structures) to find a book quickly. They might say, "Put all books alphabetically on the shelves" or "Group them by color and size." These rules work well for everyone, but they don't know your specific habits. Maybe you always borrow mystery novels, or maybe your library has a weird pattern where 90% of the books are about cats.

This paper asks a bold question: Can we teach a computer to invent its own library filing system from scratch, just by looking at the books and practicing how to find them?

The authors say yes. They created a "learning machine" that doesn't just follow rules; it discovers the rules itself.

The Two-Part Team

The system they built is like a team of two robots working together:

The Organizer (Data-Processing Network): This robot looks at the messy pile of data (the books) and figures out the best way to rearrange them. It doesn't just sort them alphabetically; it learns to sort them in a way that makes the next robot's job easier.
The Searcher (Query-Execution Network): This robot is given a specific question (e.g., "Find the book about cats"). It is only allowed to peek at a very small number of shelves (a limited "budget" of looks). It has to learn a strategy to find the right book as fast as possible using those few peeks.

The magic happens because they train together. The Organizer learns to arrange the books specifically to help the Searcher, and the Searcher learns how to read the Organizer's arrangement. They practice millions of times until they invent a system that works perfectly for the specific type of books they have.

What Did They Discover?

The researchers tested this on different types of "libraries" (datasets) and found that the robots reinvented famous human inventions, often improving on them:

The Simple List (1D Data): When the data was just a line of numbers, the Organizer learned to sort the numbers perfectly. The Searcher then learned a strategy better than the standard "Binary Search" (which is like guessing the middle of the list). If the numbers were usually small, the Searcher learned to start looking at the beginning of the list instead of the middle, saving time.
The 2D Map: When the data had two dimensions (like a map with X and Y coordinates), the robots learned to build a k-d tree. This is a complex way of splitting a map into smaller and smaller squares to find a location quickly. The robots figured this out without anyone telling them what a "tree" or a "split" was.
The High-Dimensional Maze: When dealing with complex data like images (which have thousands of features), the robots learned something called Locality Sensitive Hashing (LSH). Imagine taking a photo of a cat and instantly knowing it belongs in the "Cat Bucket" without looking at every other photo. The robots learned to project complex images into simple buckets, just like human experts do.
The "Heavy Hitter" Trick: In a test involving counting how often items appear (like tracking popular IP addresses on the internet), the robots learned to reserve special "VIP slots" in their memory for the most frequent items. This prevented the common items from getting mixed up with the rare ones, beating standard counting tools.

The "Aha!" Moment

The most surprising part is that the robots didn't need a human to say, "Hey, try sorting this!" or "Use a tree structure!" They started with random noise and, through trial and error, reverse-engineered these classic computer science algorithms on their own.

In one experiment with images of numbers, the robots learned to recognize that the images were actually numbers, sorted them by value, and then searched them efficiently—all without being told what a "number" was or how to sort them. They just learned that "similar-looking images" should be grouped together to make searching faster.

The Catch (Limitations)

The paper is honest about its limits:

Scale: The experiments were done on relatively small libraries (around 100 to 500 items). Real-world libraries have millions. The robots might get overwhelmed with that much data right now.
Speed: The robots take a long time to "think" (pre-process) before they can start searching. In real life, we often need instant answers.
Black Box: While the robots found great solutions, we don't always have a simple mathematical proof explaining why their specific arrangement works. We just know it works because we tested it.

The Bottom Line

This paper proves that neural networks can act like algorithmic inventors. Instead of humans designing the filing system, we can let the computer discover the most efficient way to organize and search data based on the specific patterns of the data it sees. It's like giving a robot a messy room and a limited amount of time to find a specific toy, and watching it invent a new way to organize the room that is even better than what a human would have designed.

Technical Summary: Discovering Data Structures End-to-End

Problem Statement

The paper investigates whether neural networks can discover data structures and query algorithms from scratch, without human-designed priors or careful initialization. The primary focus is the Nearest Neighbor (NN) search problem, where given a dataset $D$ and a query $q$ , the goal is to find $y = \arg\min_{x \in D} \text{dist}(x, q)$ . The challenge lies in framing this combinatorial problem as an end-to-end learning task that enforces constraints on space complexity (memory size) and query complexity (number of lookups), while ensuring the learned structure is efficient for specific data distributions.

A secondary application explored is frequency estimation in a streaming setting, where a model observes a stream of elements and estimates the frequency of a queried item under memory constraints.

Methodology: The Discovery Framework

The authors propose a two-component neural architecture trained end-to-end:

Data-Processing Network:
- Role: Transforms a raw dataset $D$ into a structured database $\hat{D}$ .
- Architecture: Primarily an 8-layer Transformer (based on NanoGPT) with quadratic attention, though linear attention variants are explored for efficiency.
- Mechanism: Instead of outputting the transformed data directly, the network learns a ranking (scalar value) for each data point. A differentiable sorting function then reorders the dataset based on these rankings to produce $\hat{D}$ . This permutation-based approach avoids the difficulty of learning to reproduce raw data values while preserving the ability to represent any ordering.
- Extension: In scenarios allowing extra space, the network can output additional pre-computed vectors or statistics to aid the query phase.
Query-Execution Network:
- Role: Retrieves the answer to a query $q$ by performing a fixed budget of $M$ lookups into $\hat{D}$ .
- Architecture: A sequence of $M$ independent Multi-Layer Perceptrons (MLPs), $Q_1, \dots, Q_M$ .
- Mechanism: Each query model $Q_i$ takes the query $q$ and the history of previous lookups $H_i = \{(m_1, v_1), \dots, (m_{i-1}, v_{i-1})\}$ as input. It outputs a one-hot vector $m_i$ indicating the position to access in $\hat{D}$ .
- Differentiability: To handle the discrete nature of memory access during training, the model uses Gumbel-Softmax with noise injection on the logits. This encourages sparsity and robustness, allowing the model to learn discrete selection strategies via gradient descent. At inference, a hardmax is applied.

Training Objective: The system is trained on tuples $(D, q, y)$ where $y$ is the ground-truth nearest neighbor. The loss is minimized between the predicted output $\hat{y}$ and $y$ (using squared loss or cross-entropy depending on the setting). The training is fully end-to-end, optimizing both the data structure construction and the query algorithm simultaneously.

Key Contributions and Results

1. Recovery of Classical Algorithms

The framework successfully discovers known optimal algorithms without explicit supervision:

1D NN Search:
- On uniform distributions, the model learns to sort the data and performs a search strategy superior to binary search, resembling interpolation search by leveraging distributional knowledge to start closer to the target.
- On "hard" distributions (where query position offers no prior), the model discovers binary search, achieving worst-case optimality.
- On Zipfian distributions, the model adapts to the skew, outperforming standard binary search.
2D NN Search:
- On uniform distributions, the model learns a permutation that groups spatially close points together, outperforming standard k-d trees.
- On hard distributions, the model learns a recursive partitioning strategy that closely mimics the behavior of a k-d tree, splitting data along alternating dimensions.

2. High-Dimensional Representations and LSH

Hypersphere Data: On 30-dimensional unit hyperspheres, the model learns a structure resembling Locality-Sensitive Hashing (LSH). It discovers a projection onto two vectors that partitions the space, effectively mapping similar points to the same "buckets" for efficient retrieval.
MNIST (3-digit): When tasked with finding the nearest neighbor based on numerical label (rather than pixel distance), the model learns to:
1. Extract 1D features from images that capture numerical order.
2. Sort these features.
3. Perform an interpolation-like search.
  This is achieved entirely from scratch without explicit supervision on the numerical labels or the metric space.
Real-World Benchmarks: The model performs competitively with hand-designed Learning-to-Hash baselines (ITQ, K-Means, NeuralLSH) on FashionMNIST and SIFT datasets.

3. Frequency Estimation in Streams

Applied to a streaming frequency estimation problem (Zipfian distribution), the model outperforms the standard CountMinSketch (CMS) algorithm.
Insight Transfer: The learned model utilizes a smaller update delta (increment value) than the standard CMS (which uses $\Delta=1$ ). The authors use this insight to design an augmented CMS with a custom update delta, which improves performance on real-world CAIDA IP traffic datasets by up to $1.9\times$ .

4. Utilization of Extra Space

The framework demonstrates the ability to learn how to use additional memory. In 1D and 30D settings, the data-processing network learns to pre-compute and store statistics (e.g., weighted sums of sorted inputs) that allow the query network to find the nearest neighbor with fewer lookups.

Significance and Claims

The paper claims to take a first step toward end-to-end data structure learning. Its primary significance lies in:

Feasibility: Demonstrating that deep learning models can discover meaningful, classical data structures (sorting, binary search, k-d trees, LSH) and query algorithms from scratch.
Distribution Adaptivity: Showing that learned structures can adapt to specific data distributions (e.g., uniform vs. hard vs. Zipfian) to outperform worst-case optimal hand-designed structures.
Unified Framework: Proposing a general paradigm for framing data structure problems as end-to-end learning tasks with space and query constraints, applicable beyond NN search to problems like frequency estimation, graph structures, and sparse matrix compression.
Insight Generation: Highlighting that small-scale learned models can yield transferable insights for improving classical algorithms (as seen in the augmented CountMinSketch).

Limitations and Scope

The authors explicitly state several limitations:

Scale: Experiments are limited to small datasets ( $N \approx 100$ , up to $N=500$ ). Scaling to practical, large-scale datasets remains an open challenge due to computational costs (quadratic attention, memory for large output vectors).
Efficiency: The study focuses on space and query complexity. It does not fully optimize for preprocessing time or neural network inference costs, which are critical for practical deployment.
Guarantees: Unlike learning-augmented algorithms that offer provable guarantees, end-to-end learned structures lack theoretical bounds and interpretability.
Failure Modes: The paper notes that learning can fail in specific settings (e.g., high-skew Zipfian distributions with randomized rank orders) where careful initialization or inductive biases might be required.

The paper concludes that while current models are not yet ready for production-scale replacement of classical data structures, they serve as powerful tools for exploring the algorithmic design space and discovering new heuristics.