Learning Partition Trees for Nearest Neighbor Search

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.

Prune, Don't Rebuild: Efficiently Tuning $α$-Reachable Graphs for Nearest Neighbor Search

Vector similarity search is an essential primitive in modern AI and ML applications. Most vector databases adopt graph-based approximate nearest neighbor (ANN) search algorithms, such as DiskANN (Subramanya et al., 2019), which have demonstrated state-of-the-art empirical performance. DiskANN's graph construction is governed by a reachability parameter $α$, which gives a trade-off between construction time, query time, and accuracy. However, adaptively tuning this trade-off typically requires rebuilding the index for different $α$ values, which is prohibitive at scale. In this work, we propose RP-Tuning, an efficient post-hoc routine, based on DiskANN's pruning step, to adjust the $α$ parameter without reconstructing the full index. Within the $α$-reachability framework of prior theoretical works (Indyk and Xu, 2023; Gollapudi et al., 2025), we prove that pruning an initially $α$-reachable graph with RP-Tuning preserves worst-case reachability guarantees in general metrics and improved guarantees in Euclidean metrics. Empirically, we show that RP-Tuning accelerates DiskANN tuning on four public datasets by up to $43\times$ with negligible overhead.