Jointly Optimizing Query Encoder and Product Quantization to Improve Retrieval Performance

Abstract

Recently, Information Retrieval community has witnessed fast-paced advances in Dense Retrieval (DR), which performs first-stage retrieval with embedding-based search. Despite the impressive ranking performance, previous studies usually adopt brute-force search to acquire candidates, which is prohibitive in practical Web search scenarios due to its tremendous memory usage and time cost. To overcome these problems, vector compression methods have been adopted in many practical embedding-based retrieval applications. One of the most popular methods is Product Quantization (PQ). However, although existing vector compression methods including PQ can help improve the efficiency of DR, they incur severely decayed retrieval performance due to the separation between encoding and compression. To tackle this problem, we present JPQ, which stands for Joint optimization of query encoding and Product Quantization. It trains the query encoder and PQ index jointly in an end-to-end manner based on three optimization strategies, namely ranking-oriented loss, PQ centroid optimization, and end-to-end negative sampling. We evaluate JPQ on two publicly available retrieval benchmarks. Experimental results show that JPQ significantly outperforms popular vector compression methods. Compared with previous DR models that use brute-force search, JPQ almost matches the best retrieval performance with 30x compression on index size. The compressed index further brings 10x speedup on CPU and 2x speedup on GPU in query latency.

arXiv commentAccepted by CIKM 2021 as a full paper

References in corpus (15)

• HuggingFace's Transformers: State-of-the-art Natural Language Processing
• REALM: Retrieval-Augmented Language Model Pre-Training
• Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval
• Practical and Optimal LSH for Angular Distance
• Document Expansion by Query Prediction
• Context-Aware Sentence/Passage Term Importance Estimation For First Stage Retrieval
• Semantic Models for the First-stage Retrieval: A Comprehensive Review
• Sparse, Dense, and Attentional Representations for Text Retrieval
• RocketQA: An Optimized Training Approach to Dense Passage Retrieval for Open-Domain Question Answering
• RepBERT: Contextualized Text Embeddings for First-Stage Retrieval
• Complementing Lexical Retrieval with Semantic Residual Embedding
• Distilling Dense Representations for Ranking using Tightly-Coupled Teachers
• Overview of the TREC 2021 deep learning track
• Learning To Retrieve: How to Train a Dense Retrieval Model Effectively and Efficiently
• Differentiable Product Quantization for End-to-End Embedding Compression