Reproduce “EHI: End-to-end Learning of Hierarchical Index for Efficient Dense Retrieval” on OpenResearch

About this paper

Ramnath Kumar, Anshul Mittal, Nilesh Gupta, Aditya Kusupati, Inderjit Dhillon, Prateek Jain

Researchers at Google Inc. developed EHI, an end-to-end framework that jointly optimizes dense embedding generation and a hierarchical search index for efficient retrieval. This approach achieves accuracy comparable to exact search methods while reducing document visits by up to 80% on benchmarks like MS MARCO, significantly improving efficiency and reducing latency in large-scale search systems.

Key insights

• Jointly optimizing embedding generation and the search index structure fundamentally addresses the misalignment issue, leading to embeddings that are inherently optimized for indexing.
• Learning a hierarchical index that considers query distribution during training enables more efficient document partitioning and pruning, improving latency and precision for real-world queries.
• Dense path embeddings provide a compact, differentiable representation for index traversal, allowing for end-to-end learning of complex hierarchical search structures without requiring pre-computation or warm-starting.

Summary by alphaXiv.