Learning Balanced Tree Indexes for Large-Scale Vector Retrieval

Wuchao Li, Chao Feng, Defu Lian, Yuxin Xie, Haifeng Liu, Yong Ge, Enhong Chen

Research output: Chapter in Book/Report/Conference proceedingConference contribution


Vector retrieval focuses on finding the k-nearest neighbors from a bunch of data points, and is widely used in a diverse set of areas such as information retrieval and recommender system. The current state-of-the-art methods represented by HNSW usually generate indexes with a big memory footprint, restricting the scale of data they can handle, except resorting to a hybrid index with external storage. The space-partitioning learned indexes, which only occupy a small memory, have made great breakthroughs in recent years. However, these methods rely on a large amount of labeled data for supervised learning, so model complexity affects the generalization. To this end, we propose a lightweight learnable hierarchical space partitioning index based on a balanced K-ary tree, called BAlanced Tree Learner (BATL), where the same bucket of data points are represented by a path from the root to the corresponding leaf. Instead of mapping each query into a bucket, BATL classifies it into a sequence of branches (i.e. a path), which drastically reduces the number of classes and potentially improves generalization. BATL updates the classifier and the balanced tree in an alternating way. When updating the classifier, we innovatively leverage the sequence-to-sequence learning paradigm for learning to route each query into the ground-truth leaf on the balanced tree. Retrieval is then boiled down into a sequence (i.e. path) generation task, which can be simply achieved by beam search on the encoder-decoder. When updating a balanced tree, we apply the classifier for navigating each data point into the tree nodes layer by layer under the balance constraints. We finally evaluate BATL with several large-scale vector datasets, where the experimental results show the superiority of the proposed method to the SOTA baselines in the tradeoff among latency, accuracy, and memory cost.

Original languageEnglish (US)
Title of host publicationKDD 2023 - Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining
PublisherAssociation for Computing Machinery
Number of pages10
ISBN (Electronic)9798400701030
StatePublished - Aug 6 2023
Event29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD 2023 - Long Beach, United States
Duration: Aug 6 2023Aug 10 2023

Publication series

NameProceedings of the ACM SIGKDD International Conference on Knowledge Discovery and Data Mining


Conference29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD 2023
Country/TerritoryUnited States
CityLong Beach


  • learn to index
  • maximum inner product search (mips)
  • nearest neighbor search (nns)
  • transformer
  • tree
  • vector retrieval

ASJC Scopus subject areas

  • Software
  • Information Systems


Dive into the research topics of 'Learning Balanced Tree Indexes for Large-Scale Vector Retrieval'. Together they form a unique fingerprint.

Cite this