Matching-oriented Embedding Quantization For Ad-hoc Retrieval (2021)

Abstract: Product quantization (PQ) is a widely used technique for ad-hoc retrieval. Recent studies propose supervised PQ, where the embedding and quantization models can be jointly trained with supervised learning. However, there is a lack of appropriate formulation of the joint training objective; thus, the improvements over previous non-supervised baselines are limited in reality. In this work, we propose the Matching-oriented Product Quantization (MoPQ), where a novel objective Multinoulli Contrastive Loss (MCL) is formulated. With the minimization of MCL, we are able to maximize the matching probability of query and ground-truth key, which contributes to the optimal retrieval accuracy. Given that the exact computation of MCL is intractable due to the demand of vast contrastive samples, we further propose the Differentiable Cross-device Sampling (DCS), which significantly augments the contrastive samples for precise approximation of MCL. We conduct extensive experimental studies on four real-world datasets, whose results verify the effectiveness of MoPQ. The code is available at https://github.com/microsoft/MoPQ .

## introduction ad-hoc retrieval is critical for many intelligent services, like web search and recommender systems @xcite @xcite . given a query (e.g., a search request from user), ad-hoc retrieval selects relevant keys (e.g., webpages) from a massive set of candidates. it is usually treated as an approximate nearest neighbour search (anns) problem, where product quantization (pq) @xcite is one of the most popular solutions thanks to its competitive memory and time efficiency. pq is built upon †. work is done during the internship at microsoft. * . corresponding author. existing limitation. the original pq algorithms @xcite are non-supervised: based on the well-trained embeddings, the quantization model is learned with heuristic algorithms (e.g., k-means). in recent years, many works are dedicated to supervised pq @xcite @xcite @xcite , where the embedding model and the quantization model are trained jointly. the supervised pq requires an explicit training objective for the quantization model. most of the time, "the minimization of reconstruction loss" is used for granted: the distortion between the original embedding (z) and its quantization result (z) needs to be reduced as much as possible for all the keys (k) in database: the above objective is seemingly plausible, as a "small enough distortion" will make the quantized embeddings equally expressive as the original embeddings. however, it implicitly hypothesizes that the distortion can be made sufficiently small, which is not always realistic in practice. this is because the reconstruction loss is subject to a lower-bound determined by the codebooks scale. as over-scaled codebooks result in prohibitive memory and time costs, there will always exist a positive reconstruction loss in reality. in this case, it can be proved that the minimization of reconstruction loss doesn't lead to the optimal retrieval accuracy. it's also empirically verified that the supervised pq's advantage over the non-supervised baselines are limited and not consistently positive when the reconstruction loss minimization is taken as the training objective (see section 2.2 and 4 for detailed analysis). our work. to address the above challenge, we propose the matching-oriented product quantization (mopq), with a novel objective mcl formulated to optimize pq's retrieval accuracy, together with a sample augmentation strategy dcs to ensure the effective minimization of mcl. in summary, our work identifies a long-existing defect about the training objective of supervised pq; meanwhile, a simple but effective remedy is proposed, which optimizes the expected retrieval accuracy of pq. we make extensive experimental studies with four benchmark text retrieval datasets, where our proposed methods significantly outperform the sota supervised pq baselines. our code and datasets will be made public-available to facilitate the research progress in related areas. ## revisit of supervised pq we start with the preliminaries of pq's application in ad-hoc retrieval. then, we analyze the defect of taking the reconstruction loss minimization as supervised pq's training objective. ## matching-oriented pq we present the matching-oriented pq (mopq) in this section. in mopq, a new quantization objective mcl is proposed, whose minimization optimizes the query-key's matching probability, therefore bringing about the optimal retrieval accuracy. besides, the dcs method is introduced, which facilitates the effective minimization of mcl. ## conclusion in this paper, we propose mopq to optimize pq's ad-hoc retrieval accuracy. a systematic revisit is made for the existing supervised pq, where we identify the limitation of using reconstruction loss minimization as the training objective. we propose mcl as our new training objective, where the model can be learned to maximize the querykey matching probability to achieve the optimal retrieval accuracy. we further leverage dcs for contrastive sample argumentation, which ensures the effective minimization of mcl. the experiment results on 4 real-world datasets validate the effectiveness of our proposed methods. the z i will still be mapped to j-th codeword in codebook c i . in other words, the original codeword assignment will not be affected by i .