Learning Distance-Aware Space Partitions for Approximate Nearest Neighbor Search

Approximate Nearest Neighbor Search (ANNS) is a fundamental problem widely applied in information retrieval and data mining, with graph-based methods gaining particular interest due to their outstanding efficiency and query accuracy. However, as the scale of high-dimensional data continues to grow, efficient space partitioning becomes essential for supporting distributed vector data processing. We model space partitioning as a balanced graph partitioning problem, but traditional graph partitioning algorithms, which primarily aim to minimize edge cuts, may inadvertently separate closer neighbors. To overcome this limitation, we propose a distance-aware multilevel partitioning algorithm that combines distance-aware coarsening and refinement strategies, effectively reducing the incidence of closer nodes being cut. We train a lightweight model to predict partition where the neighbors of a query point are most likely to be located by minimizing the Kullback-Leibler (KL) divergence between the predicted partition distribution and the ground truth distribution derived from neighbor counts. Additionally, our approach dynamically repartitions boundary nodes after each training iteration, enabling alternating updates that optimize both the classifier and the quality of the partitioning. Experimental results demonstrate the superiority of the proposed method to the baselines in the tradeoff among search efficiency and partitioning quality.