EHI: End-to-end Learning of Hierarchical Index for Efficient Dense Retrieval

TL;DR

End-to-end Hierarchical Indexing (EHI) is proposed, a novel method that directly addresses this issue by jointly optimizing embedding generation and ANNS structure by jointly optimizing embedding generation and ANNS structure.

Abstract

Dense embedding-based retrieval is widely used for semantic search and ranking. However, conventional two-stage approaches, involving contrastive embedding learning followed by approximate nearest neighbor search (ANNS), can suffer from misalignment between these stages. This mismatch degrades retrieval performance. We propose End-to-end Hierarchical Indexing (EHI), a novel method that directly addresses this issue by jointly optimizing embedding generation and ANNS structure. EHI leverages a dual encoder for embedding queries and documents while simultaneously learning an inverted file index (IVF)-style tree structure. To facilitate the effective learning of this discrete structure, EHI introduces dense path embeddings that encodes the path traversed by queries and documents within the tree. Extensive evaluations on standard benchmarks, including MS MARCO (Dev set) and TREC DL19, demonstrate EHI's superiority over traditional ANNS index. Under the same computational constraints, EHI outperforms existing state-of-the-art methods by +1.45% in MRR@10 on MS MARCO (Dev) and +8.2% in nDCG@10 on TREC DL19, highlighting the benefits of our end-to-end approach.

Figures (15)

• Figure 1: Motivation of EHI.
• Figure 2: Comparison between EHI's end-to-end training paradigm in comparison to other approaches such as Dual encoder + ANNS and Generative Retrieval approaches which require a multi-stage training.We denote stages which are not optimized for the task objective by a grey shade. Note that unlike DE + ANNS, and Generative Retrieval, ${\rm EHI}$ optimizes for recall in both encoder and indexer.
• Figure 3: ${\rm EHI}$ is an end-to-end hierarchical indexer which comprises an encoder and a hierarchical tree as the indexer where the entire pipeline is learnable and differentiable. Here, variables $V_{98}$, $V_{123}$, and $V_{576}$ are dense representations (embeddings) of the text and $P_{98}$, $P_{123}$, and $P_{576}$ are path embeddings of the respective samples. To efficiently train ${\rm EHI}$ without any warm starting, we use a combination of objectives - $L_{\text{siamese}}$, $L_{\text{indexing}}$, $L_{\text{intra-leaf}}$ (see Section \ref{['sec:algo']} for details).
• Figure 4: ${\rm EHI}$ is significantly more accurate than DE + ScaNN or Faiss-IVF, especially when restricted to visit a small fraction of documents. See Figure \ref{['fig:scifact_fiqa_results_appendix']} in Appendix for results on Scifact, Fiqa.
• Figure 5: Ablations studies depicting the various properties of ${\rm EHI}$'s training paradigm.