Efficient Document Retrieval by End-to-End Refining and Quantizing BERT Embedding with Contrastive Product Quantization (2022)

Abstract: Efficient document retrieval heavily relies on the technique of semantic hashing, which learns a binary code for every document and employs Hamming distance to evaluate document distances. However, existing semantic hashing methods are mostly established on outdated TFIDF features, which obviously do not contain lots of important semantic information about documents. Furthermore, the Hamming distance can only be equal to one of several integer values, significantly limiting its representational ability for document distances. To address these issues, in this paper, we propose to leverage BERT embeddings to perform efficient retrieval based on the product quantization technique, which will assign for every document a real-valued codeword from the codebook, instead of a binary code as in semantic hashing. Specifically, we first transform the original BERT embeddings via a learnable mapping and feed the transformed embedding into a probabilistic product quantization module to output the assigned codeword. The refining and quantizing modules can be optimized in an end-to-end manner by minimizing the probabilistic contrastive loss. A mutual information maximization based method is further proposed to improve the representativeness of codewords, so that documents can be quantized more accurately. Extensive experiments conducted on three benchmarks demonstrate that our proposed method significantly outperforms current state-of-the-art baselines.

## introduction in the era of big data, approximate nearest neighbor (ann) search has attracted tremendous attention thanks to its high search efficiency and extraordinary performance in modern information retrieval systems. by quantizing each document as a compact binary code, semantic hashing @xcite has become the main solution to ann search due to the extremely low cost of calculating hamming distance between binary codes. one of the main approaches for unsupervised semantic hashing methods is established on generative models @xcite @xcite , which encourage the binary codes to be able to reconstruct the input documents. alternatively, some other methods are driven by graphs @xcite @xcite , hoping the binary codes can recover the neighborhood relationship. though these methods have obtained great retrieval performance, there still exist two main problems. firstly, these methods are mostly established on top of the outdated tfidf features, which do not contain various kinds of important information of documents, like word order, contextual information, etc. in recent years, pre-trained language models like bert have achieved tremendous success in various downstream tasks. thus, a natural question to ask is whether we can establish efficient retrieval methods on bert embeddings. however, it has been widely reported that bert embeddings are not suitable for semantic similarity-related tasks @xcite , which perform even worse than the traditional glove embeddings @xcite . @xcite attribute this to the "anisotropy" phenomenon that bert embeddings only occupy a narrow cone in the vector space, causing the semantic information hidden in bert embeddings not easy to be leveraged directly. thus, it is important to investigate how to effectively leverage the bert embeddings for efficient document retrieval. secondly, to guarantee the efficiency of retrieval, most existing methods quantize every document to a binary code via semantic hashing. there is no doubt that the hamming distance can improve the retrieval efficiency significantly, but it also restricts the representation of document similarities seriously since it can only be an integer from -b to b for b-bit codes. recently, an alternative approach named product quantization @xcite @xcite has been proposed in the computer vision community to alleviate this problem. basically, it seeks to quantize every item to one of the codewords in a codebook, which is represented by a cartesian product of multiple small codebooks. it has been shown that product quantization is able to deliver superior performance than semantic hashing while keeping the cost of computation and storage relatively unchanged. however, this technique has rarely been explored in unsupervised document retrieval. motivated by the two problems above, in this paper, we propose an end-to-end contrastive product quantization model to jointly refine the original bert embeddings and quantize the refined embeddings into codewords. specifically, we first transform the original bert embeddings via a learnable mapping and feed the transformed embedding into a probabilistic product quantization module to output a quantized representation (codeword). to preserve as much semantic information as possible in the quantized representations, inspired by recent successes of contrastive learning, a probabilistic contrastive loss is designed and trained in an end-to-end manner, simultaneously achieving the optimization of refining and quantizing modules. later, to further improve the retrieval performance, inspired by the recent development of clustering, a mutual information maximization based method is further developed to increase the representativeness of learned codewords. by doing so, the cluster structure hidden in the dataset of documents could be kept soundly, making the documents be quantized more accurately. extensive experiments are conducted on three real-world datasets, and the experimental results demonstrate that our proposed method significantly outperforms current state-ofthe-art baselines. empirical analyses also demonstrate the effectiveness of every proposed component in our proposed model. ## preliminaries of product quantization for information retrieval in fields of efficient information retrieval, a prevalent approach is semantic hashing, which maps every item x to a binary code b and then uses hamming distances to reflect the semantic similarity of items. thanks to the low cost of computing hamming distance, the retrieval can be performed very efficiently. however, the hamming distance can only be an integer from -b to b for b-bit codes, which is too restrictive to reflect the rich similarity information. an alternative approach is vector quantization (vq) @xcite , which assigns every item with a codeword from a codebook c. the codeword in vq could be any vector in r d , rather than limited to the binary form as in semantic hashing. by storing pre-computed distances between any two codewords in a table, the distance between items can be obtained efficiently by looking up the table. however, to ensure competitive performance, the number of codewords in a codebook needs to be very large. for example, there are 2 64 different codes for a 64-bit binary code, and thus the number of codewords in vq should also be of this scale, which however is too large to be handled. to tackle this issue, product quantization @xcite proposes to represent the codebook c as a cartesian product of m small codebooks ## the end-to-end joint refining and quantizing framework to retrieve semantic-similar documents, a core problem is how to produce for every document a quantized representation that preserves the semantic-similarity information of original documents. in this section, a simple two-stage method is first proposed, and then we develop an end-to-end method that is able to directly output representations with desired properties. ## related work as the main solution to efficient document retrieval, unsupervised semantic hashing has been studied for years. many existing unsupervised hashing methods are established on the generative models encouraging the binary codes to reconstruct the original document. for example, vdsh @xcite proposes a two-stage scheme, in which it first learns the continuous representations under the vae (kingma and welling, 2014) framework, and then cast them into the binary codes. to tackle the two-stage training issue, nash @xcite presents an end-to-end generative hashing framework where the binary codes are treated as bernoulli latent variables, and introduces the straight-through @xcite estimator to estimate the gradient w.r.t. the discrete variables. @xcite employs the mixture priors to empower the binary code with stronger expressiveness, therefore resulting in better performance. further, corrsh @xcite employs the distribution of the boltzmann machine to introduce correlations among the bits of binary codes. @xcite proposes the auxiliary implicit topic vectors to address the issue of information loss in the few-bit scenario. also, a handful of recent works focus on how to inject the neighborhood information of the graph under the vae framework. @xcite ; @xcite seeks to learn the optimal binary codes that can reconstruct neighbors of original documents. a ranking loss is introduced in @xcite to accurately characterize the correlation between documents. @xcite first proposes to integrate the semantics and neighborhood information with a graph-driven generative model. beyond generative hashing methods, studies on hashing via the mutual information (mi) principle emerges recently. ammi (stratos and wiseman, 2020) learns a high-quality binary code by maximizing the mi between documents and binary codes. dhim @xcite first compresses the bert embeddings into binary codes by maximizing the mi between global codes and local codes from documents. another efficient retrieval mechanism is product quantization. the product quantization (pq) @xcite and its improved variants such as optimized pq @xcite and locally optimized pq @xcite are proposed to retrain richer distance information than hashing methods while conducting the retrieval efficiently. these shallow unsupervised pq methods are often based on the well-trained representation and learn the quantization module with heuristic algorithms @xcite , which often can not achieve satisfactory performance. in this paper, we propose an unsupervised mi-improved end-to-end unsupervised product quantization model micpq. we notice that the proposed micpq is somewhat similar to recent works w.r.t pq @xcite in the computer vision community. however, @xcite focuses on analyzing the performance difference between different forms of contrastive losses, while we focus on how to design an end-to-end model to jointly refine and quantize the bert embedding via contrastive product quantization from a probabilistic perspective. @xcite concentrates on how to improve the codeword diversity to prevent model degeneration by regularizer design; whereas there is no "model degeneration" phenomenon observed in our model, and we use the mutual information maximization based method to increase the representativeness of codewords to further improve the retrieval performance. ## conclusion in this paper, we proposed a novel unsupervised product quantization model, namely micpq. in micpq, we managed to develop an end-toend probabilistic contrastive product quantization model to jointly refine the original bert embeddings and quantize the refined embeddings into codewords, with a probabilistic contrastive loss designed to preserve the semantic information in the quantized representations. moreover, to improve the representativeness of codewords for keeping the cluster structure of documents soundly, we proposed to maximize the mutual information between data and the codeword assignment. extensive experiments showed that our model significantly outperformed existing unsupervised hashing methods.