Learning to Design Data-structures: A Case Study of Nearest Neighbor Search

Omar Salemohamed, Vatsal Sharan, Shivam Garg, Laurent Charlin, Gregory Valiant

Keywords: sorting, nearest-neighbors, sparsity, data-structures

TL;DR: We propose a model for learning data-structures from scratch and apply it to the nearest neighbor problem.

Abstract: We propose a general framework for automating data-structure design and apply it to the problem of nearest neighbor search. Our model adapts to the underlying data distribution and provides fine-grained control over query and space complexity, enabling the discovery of solutions tailored to problem-specific constraints. We are able to reverse-engineer learned algorithms in several settings. In 1D, 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, have elements of locality-sensitive hashing.

Submission Number: 21