Research map / Learning theory and algorithms / Prediction and decision losses
Discovering Data Structures: Nearest Neighbor Search and Beyond
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes learning data structures end-to-end from scratch with a transformer data-processing network and per-lookup MLP query networks, and studies nearest-neighbor search (1D, 2D, high-D) and frequency estimation.
Problem
Whether neural networks can discover data structures and query algorithms end-to-end, under explicit space and query-complexity constraints, without hand-designed initialization, focusing on nearest-neighbor search.
Approach
Jointly train a transformer that ranks/permutes the dataset into a data structure and M MLP query models that perform adaptive discrete lookups (softmax with logit noise during training, hardmax at inference), trained from scratch on sampled (dataset, query, nearest neighbor) tuples so the structure and query algorithm co-adapt.
Key results
- On 1D uniform NN search with N=100, M=7, the learned E2E model outperforms binary search and learns sorting with ~99.5% correct positions.
- On 2D uniform data, E2E reaches ~75% accuracy vs ~52% for k-d tree (N=100, M=6).
- On 3-digit MNIST, E2E achieves 97.7% NN accuracy vs 76.2% for an oracle-bucket LSH baseline; on hypersphere (d=30) E2E gets 35.0 vs LSH 30.0.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use E2E instead of LSH | High-dimensional nearest-neighbor search · Hypersphere | accuracy: 30.0 -> 35.0 | Helped |
| Use E2E instead of NeuralLSH | High-dimensional nearest-neighbor search · FashionMNIST | accuracy: 68.2 -> 67.2 | Hurt |
| Use E2E instead of KMeans | High-dimensional nearest-neighbor search · SIFT | accuracy: 47.1 -> 46.9 | Hurt |
| Use E2E instead of LSH | High-dimensional nearest-neighbor search · MNIST (3-digit) | accuracy: 76.2 -> 97.7 | Helped |
| Use E2E instead of a k-d tree | 2D nearest-neighbor search · 2D uniform | accuracy: ~52% -> ~75% | Helped |
| Increase learned extra space from 0 to 128 vectors | E2E · 1D uniform | accuracy: 44.0 -> 91.3 | Helped |
| Learn data rankings rather than freeze the data-processing network (ablation) | E2E · 1D uniform | accuracy | Helped |
| Add query-history conditioning (ablation) | E2E · 1D uniform | nearest-neighbor accuracy | Helped |
| Keep a permutation-based data-processing network instead of directly transforming data points (ablation) | E2E · 1D uniform | nearest-neighbor accuracy | Helped |
| Use a learned smaller update delta | CountMinSketch · CAIDA IP dataset | relative performance: ≈1.9× | Helped |
Surprising findings
- The model discovers binary-search-like and interpolation-like search, and k-d-tree-like and LSH-like partitions without supervision on these primitives.
- Learned frequency-estimation structure uses a smaller update delta than CountMinSketch, and this insight yields an augmented CMS up to ~1.9x better on CAIDA IP data.
Limitations
- Experiments limited to small scale (mostly N=100, some N=500); scaling remains open.
- Learned structures lack provable guarantees and interpretability; preprocessing and inference cost not fully addressed.
- End-to-end training fails in some cases and hard distributions are difficult to train for N>16.
Experimental setup
Datasets: synthetic uniform (-1,1) 1D and 2D synthetic hard distribution synthetic Zipfian uniform unit hypersphere 30D 3-digit MNIST FashionMNIST SIFT CAIDA IP trace
Benchmarks: synthetic 1D/2D/30D NN search FashionMNIST (PCA 100D) SIFT 3-digit MNIST NN CAIDA IP frequency estimation Zipfian frequency estimation
Metrics: nearest neighbor accuracy mean squared error mean absolute error of frequency estimates
Baselines: binary search k-d tree LSH ITQ KMeans NeuralLSH CountMinSketch learning-augmented treap
Builds on: NanoGPT Gumbel softmax k-d tree SimHash CountMinSketch
Closest papers by meaning
- Learning-Augmented Search Data Structures (ICLR 2025 · similarity 0.88)
- Learning Space Partitions for Nearest Neighbor Search (ICLR 2020 · similarity 0.88)
- Graph-Based Algorithms for Diverse Similarity Search (ICML 2025 · similarity 0.86)
- Approximate Nearest Neighbor Search with Window Filters (ICML 2024 · similarity 0.85)
- Revisiting Nearest Neighbor for Tabular Data: A Deep Tabular Baseline Two Decades Later (ICLR 2025 · similarity 0.85)
- Distance Adaptive Beam Search for Provably Accurate Graph-Based Nearest Neighbor Search (NeurIPS 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.