Learning Partition Trees for Nearest Neighbor Search

2026-07-10

AI Summary

This work addresses the design of nearest neighbor search data structures tailored to a given query distribution. It proposes the first algorithm capable of efficiently learning an approximately optimal balanced halfspace partitioning tree under Gaussian-like distributional assumptions. By formulating tree construction as a balanced halfspace cut problem and integrating polynomial threshold functions with a distribution-aware learning strategy, the method circumvents the NP-hard regularized optimization typically involved. Under the assumption that a perfect partitioning tree exists, the approach achieves query time better than $O(nd)$ while ensuring provably bounded cutting error in the learned tree, thereby significantly enhancing both the efficiency and theoretical guarantees of data-driven nearest neighbor search.

Abstract

We study nearest neighbor search from the perspective of data-driven algorithm design: given a dataset $P \subset \mathbb{R}^d$ of size $n$ and sample access to a query distribution over $\mathbb{R}^d$, the goal is to learn a data structure optimized for queries drawn from that specific distribution. We focus on the class of balanced halfspace trees, which naturally abstracts space-partitioning frameworks like locality-sensitive hashing. Assuming Gaussian-like marginal conditions on the dataset and query distribution, we give an efficient algorithm that learns a tree achieving $o(nd)$ query time, provided that a perfect tree exists. At the core of our algorithmic approach is the balanced halfspace cut problem, where we are given a distribution over $\mathbb{R}^d \times \mathbb{R}^d$ and must find a balanced halfspace that minimizes the fraction of cut pairs. We prove that without distributional assumptions, finding the optimal balanced halfspace is NP-hard. To circumvent this computational barrier, we design an efficient improper learning algorithm: if the optimal halfspace cuts an $α$ fraction of pairs, our algorithm outputs a balanced polynomial threshold function of degree $\tilde{O}(1/\varepsilon^2)$ that cuts at most an $O(\sqrt{α+\varepsilon})$ fraction.

Problem

• nearest neighbor search
• data-driven algorithm design
• balanced halfspace trees
• query distribution
• space-partitioning

Innovation

• data-driven nearest neighbor search
• balanced halfspace trees
• improper learning
• polynomial threshold functions
• distribution-aware indexing