Improving Document Representations by Generating Pseudo Query Embeddings for Dense Retrieval (2021)

Abstract: Recently, the retrieval models based on dense representations have been gradually applied in the first stage of the document retrieval tasks, showing better performance than traditional sparse vector space models. To obtain high efficiency, the basic structure of these models is Bi-encoder in most cases. However, this simple structure may cause serious information loss during the encoding of documents since the queries are agnostic. To address this problem, we design a method to mimic the queries to each of the documents by an iterative clustering process and represent the documents by multiple pseudo queries (i.e., the cluster centroids). To boost the retrieval process using approximate nearest neighbor search library, we also optimize the matching function with a two-step score calculation procedure. Experimental results on several popular ranking and QA datasets show that our model can achieve state-of-the-art results while still remaining high efficiency.

## introduction given a query and a collection of documents, the document retrieval task is to rank the documents based on their relevance with the query. to retrieve the target documents efficiently, most existing work adopts a two-stage fashion which retrieves a subset of candidate documents from the whole corpus by a recall model and then re-rank them by a sophisticated ranking model. in the first stage, many approaches use traditional information retrieval methods including bm25 based on sparse bag-of-word representation. since the recall of the first-stage model determines the upper bound of the ranking quality, there is lots of work focusing on improving the recall performance @xcite @xcite . in contrast to the sparse representations, dense representations encoding semantic information can enhance the retrieval performance by overcoming the limitations like term mismatching. they are usually produced by neural encoders whose parameters are learnable. recently, inspired by the great success of pre-trained language models like bert/roberta @xcite in nlp applications, the dense passage retriever is proposed which encodes the documents by fine-tuning the huge language models @xcite and achieves state-of-theart results benefiting from their powerful contextual semantic representative ability. following the typical fine-tuning paradigm on many nlp tasks @xcite , a bert encoder usually takes the concatenation of the query and document text as input and performs a full selfattention across the input tokens. such architecture is called cross-encoder @xcite . although it can achieve better performance than other architectures, it is infeasible in the recall stage since it needs to recompute the representation of each document in the corpus once a new query is provided. in contrast, bi-encoder @xcite encodes the queries and documents separately and computes the matching scores between their dense representations. since the documents in the corpus keep unchanged most of the time, the representation of the documents can be stored in advance for future use. with the help of approximate nearest neighbor (ann) search approaches @xcite , the retrieval process can be further boosted. although gaining retrieval efficiency, bi-encoder sacrifices retrieval accuracy comparing to the cross-encoder. to enrich the representations of the documents produced by bi-encoder, some researchers extend the original bi-encoder by employing more delicate structures like laterinteraction(khattab and zaharia, 2020), poly-encoder @xcite , multi-embedding model @xcite . increasing a little computational overhead, these models can gain much improvement of the encoding quality while remaining the fast retrieval characteristic of bi-encoder. similar to these models, we focus on improving the effectiveness of bi-encoder. in this work, we think that the limitation of the bi-encoder origins from the agnostic nature of query when encoding the documents independently, i.e., the encoder cannot know what query could be potentially answered by the input document. as it is very common that a document with hundreds of tokens contains several distinct topics, some important semantic information might be easily missed or biased by each other without knowing the query. to alleviate the query agnostic problem, we propose a novel approach that mimics multiple potential queries corresponding to the input document and we call them "pseudo query embeddings". ideally, each of the pseudo query embeddings corresponds to a semantic salient fragment in the document which is similar to a semantic cluster of the document. thus, we implement the process by a clustering algorithm (i.e., k-means in this work) and regard the cluster centroids as the pseudo query embeddings. we generate and store all of the embeddings in an offline manner, thereby not only improving the encoding quality but also remaining the online computation unchanged. during the inference, the multiple pseudo query embeddings should be first aggregated through a softmax function and then the relevance score with the query embedding is computed. unfortunately, directly applying softmax aggregation is not supported in the existing ann search library. thus, we first filter some documents in which all of the embeddings have low relevance scores and then perform the whole aggregation and score function using the filtered embeddings. our main contributions can be summarized as follows: ## related work in this section, we will review the existing work related with the first-stage retrieval. according to the representations of text, the first stage retrieval approaches can be classified into two categories. one is based on the high-dimensional sparse representation and the other is based on the low-dimensional continuous representation. traditional sparse vector space models weight the terms by their frequency information. in last few years, some researchers intend to weight the document and query terms adaptively by a neural network which could leverage some semantical information @xcite . recently, a trend of leveraging the deep pre-trained language models to weight or augment the document/query terms is emerged. deepct @xcite uses bert to learn the term importance and weight all of the terms. doct5query @xcite ) augments the document with possible query terms which are generated by a sequence-to-sequence model. in contrast, the dense retrieval approaches map the text to continuous vectors which are mostly generated by neural networks. models like dssm @xcite ,clsm @xcite , desm @xcite encode the query and document using their n-gram features or word embeddings independently and then compute their similarities. recently, the dense retrieval approaches also tend to make use of the pre-trained language models. sentence-bert @xcite is a typical bi-encoder model which encodes the text using bert and calculates the similarity scores by the combination of several basic operations. inspired by the interaction-based neural re-rankers, khattab and zaharia(2020) propose a later-interaction mechanism. later on, some variants @xcite are proposed. @xcite identify that the negative samples during training may not be representative, lowering the training difficulty. therefore, they propose a model to construct hard negative samples dynamically during training. comparing to existing work, our work serves the first stage of document retrieval and presents a new method to generate document representations which borrows the clustering technique to generate pseudo query embeddings from documents. ## dense document retrieval in this section, we introduce the original biencoder architecture and several existing variants. then, we present our model in detail and describe the similarities and differences between our model and those bi-encoder variants. ## conclusions in this paper, we propose a method to improve the performance of the first-stage retrieval model which is based on bi-encoder and semi-interactive aggregator. specifically, our method mimics the real queries by an iterative k-means clustering algorithm. to accelerate the retrieval process, we also optimize the softmax matching function by filtering out some documents using argmax operation. we conduct experiments on the ms marco and openqa datasets. through the analysis of the retrieval quality and efficiency, we can confirm the proposed approach is both effective and efficient.