Jointly Optimizing Query Encoder and Product Quantization to Improve Retrieval Performance

TL;DR

JPQ introduces a joint optimization framework for query encoding and product quantization, significantly enhancing retrieval efficiency and performance in dense retrieval systems by reducing index size and speeding up query processing.

Contribution

It proposes an end-to-end training method for query encoder and PQ, improving retrieval accuracy while achieving high compression and speedup.

Findings

JPQ outperforms existing vector compression methods in retrieval accuracy.

JPQ achieves nearly the same performance as brute-force search with 30x smaller index.

JPQ provides 10x CPU and 2x GPU speedup in query latency.

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.