Learning Space Partitions for Nearest Neighbor Search
This work addresses the need for efficient nearest neighbor search algorithms, which is a fundamental problem in machine learning and data retrieval, with incremental improvements over prior methods.
The paper tackles the problem of fast nearest neighbor search by developing a new framework for building space partitions, which reduces the problem to balanced graph partitioning and supervised classification, resulting in Neural LSH that consistently outperforms existing methods on standard benchmarks.
Space partitions of $\mathbb{R}^d$ underlie a vast and important class of fast nearest neighbor search (NNS) algorithms. Inspired by recent theoretical work on NNS for general metric spaces [Andoni, Naor, Nikolov, Razenshteyn, Waingarten STOC 2018, FOCS 2018], we develop a new framework for building space partitions reducing the problem to balanced graph partitioning followed by supervised classification. We instantiate this general approach with the KaHIP graph partitioner [Sanders, Schulz SEA 2013] and neural networks, respectively, to obtain a new partitioning procedure called Neural Locality-Sensitive Hashing (Neural LSH). On several standard benchmarks for NNS, our experiments show that the partitions obtained by Neural LSH consistently outperform partitions found by quantization-based and tree-based methods as well as classic, data-oblivious LSH.