SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hailin Zhang $^{1,2,*}$  Yujing Wang $^{2,\dagger}$  Qi Chen $^{2}$  Ruiheng Chang $^{2}$  Ting Zhang $^{2}$   
Ziming Miao $^{2}$  Yingyan Hou $^{2,3,*}$  Yang Ding $^{2,4,*}$  Xupeng Miao $^{5}$  Haonan Wang $^{6}$   
Bochen Pang $^{2}$  Yuefeng Zhan $^{2}$  Hao Sun $^{2}$  Weiwei Deng $^{2}$  Qi Zhang $^{2}$   
Fan Yang $^{2}$  Xing Xie $^{2}$  Mao Yang $^{2}$  Bin Cui $^{1,7,\dagger}$
$^{1}$ School of Computer Science & Key Lab of High Confidence Software Technologies, Peking University  $^{2}$ Microsoft
$^{3}$ Aerospace Information Research Institute & Key Laboratory of Target Cognition and Application Technology, Chinese Academy of Sciences
$^{4}$ Institute of Information Engineering, Chinese Academy of Sciences  $^{5}$ Carnegie Mellon University  $^{6}$ National University of Singapore
$^{7}$ National Engineering Laboratory for Big Data Analysis and Applications, Peking University  $^{1}\{z.hl, \text{bin.cui}\} @$ pku.edu.cn
$^{2}\{yujwang, cheqi, ruihengchang, tinzhan, zimiao, bopa\} @$ microsoft.com
$^{2}\{yuefzh, hasun, dedeng, qizhang, fanyang, xingx, maoyang\} @microsoft.com$ $^{3}$ houyy@aircas.ac.cn  $^{4}$ dingyang@iie.ac.cn
$^{5}$ xupeng@cmu.edu  $^{6}$ haonan.wang@u.nus.edu
Embedding-based retrieval methods construct vector indices to search for document representations that are most similar to the query representations. They are widely used in document retrieval due to low latency and decent recall performance. Recent research indicates that deep retrieval solutions offer better model quality, but are hindered by unacceptable serving latency and the inability to support document updates. In this paper, we aim to enhance the vector index with end-to-end deep generative models, leveraging the differentiable advantages of deep retrieval models while maintaining desirable serving efficiency. We propose Model-enhanced Vector Index (MEVI), a differentiable model-enhanced index empowered by a twin-tower representation model. MEVI leverages a Residual Quantization (RQ) codebook to bridge the sequence-to-sequence deep retrieval and embedding-based models. To substantially reduce the inference time, instead of decoding the unique document ids in long sequential steps, we first generate some semantic virtual cluster ids of candidate documents in a small number of steps, and then leverage the well-adapted embedding vectors to further perform a fine-grained search for the relevant documents in the candidate virtual clusters. We empirically show that our model achieves better performance on the commonly used academic benchmarks MSMARCO Passage and Natural Questions, with comparable serving latency to dense retrieval solutions.
A web search engine basically consists of two key stages, document retrieval and ranking [52, 24, 47, 9]. The document retrieval stage retrieves candidate documents related to the query, while the ranking stage re-ranks the documents using precise ranking scores. Accurate ranking models are usually deep neural networks that predict query and document relevance scores [17, 13]. To guarantee low latency
*The work was done at Microsoft.
in online systems, such ranking models can only support hundreds or thousands of documents per query, requiring the retrieval stage to recall relevant documents precisely. Therefore, improving the quality of retrieval models is of great significance to improve the effectiveness of web search engines.
Existing methods of document retrieval can be roughly divided into three categories, namely term-based, embedding-based and generation-based methods [14, 46]. Term-based methods, also called sparse retrieval, build an inverted index upon words or phrases from the whole corpus for searching [4, 56]. However, they are not aware of document semantics information and may miss target documents with different wordings. Embedding-based methods, also called dense retrieval, are proposed to leverage semantics information [52, 28]. They encode queries and documents to dense embedding vectors by a twin-tower architecture, then build a vector index and apply Approximate Nearest Neighbor (ANN) search to retrieve relevant documents for queries [12]. Such methods are widely used in real applications, since they have decent recall performance and low serving latency. Although these methods prevail in industry, they cannot fully leverage the power of deep neural networks because they are not end-to-end differentiable due to the necessity of ANN index. They suffer from the discrepancy between brute-force K-Nearest Neighbor (KNN) and ANN. The discrepancy leads to a non-trivial gap in recall performance of KNN and ANN [49].
Recently, generation-based methods [44, 46] are proposed to tackle the limitation of embedding-based retrieval solutions. Typically, a generation-based model adopts a sequence-to-sequence architecture, generating document identifiers directly based on the given queries. To achieve better results, the document identifiers need to reflect effective priors of the document semantics. For example, SEAL [2] leverages the n-grams itself as identifiers. In DSI [44] and NCI [46], documents are organized as a tree by hierarchical KMeans clustering [15], and identifiers are the codes of paths from the root to the leaf nodes. Though such methods reach better recall performance than embedding-based retrieval models on small corpus (containing less than 1 million documents), they have difficulty in scaling to larger corpus (containing more than 10 million documents), and can not be served online in industrial systems due to high latency and immutable corpus. To train a deep retrieval model that indexes larger datasets, we need much more parameters to memorize all documents semantics, and the decoder needs to compute more times by beam search at inference. Another blocking issue is that the tree index is fixed during training and inference, posing a significant challenge for adding or removing documents. As a result, the model needs to be retrained every time the corpus changes. Overall, generation-based retrieval methods are time-consuming and only applicable to static small corpus. To the best of our knowledge, they have not been successfully used in real applications with large-scale corpus due to these non-trivial challenges.
In this paper, we aim to address the aforementioned constraints of model-based methods while preserving their benefits. We propose Model-enhanced Vector Index (MEVI), which enjoys high recall performance as well as fast retrieval speed on large-scale corpus. Specifically, we construct a Residual Quantization (RQ) codebook to cluster the documents first. The RQ codebook preserves the hierarchical structure of the document clusters, which is inherently suitable for autoregressive generation. Thereafter, we build a sequence-to-sequence model to encode user queries and directly generate virtual cluster identifiers according to the RQ codebook; then we conduct efficient ANN method to search embedding vectors which are semantically relevant to the virtual cluster. During training, ground-truth data and augmented query-document pairs are fed into the model. During inference, the top K document clusters are retrieved from the RQ codebook via beam search on the decoder, and an ANN search process is conducted based on the query embedding as well as the retrieved document embeddings corresponding to the clusters. Our MEVI design solves the limitations of both traditional embedding-based methods and generation-based methods. On the one hand, we can limit the RQ codebook to a moderate size to reduce the computation time of the autoregressive decoder and thus ensure low latency. On the other hand, MEVI enables insertion or removal of documents to RQ-based index structures, such that new documents can be recalled if their virtual clusters with related semantics can be generated by the sequence-to-sequence model. By selecting a suitable size for the RQ codebook, we can balance the recall performance and the inference latency. This allows us to leverage the high efficiency of ANN and the accurate recall of deep retrieval models simultaneously. The code of MEVI is available at: https://github.com/HugoZHL/MEVI.
Our contributions are highlighted as follows.
- For the first time, we demonstrate that a novel-designed generation-enhanced model is able to handle a large corpus with millions of documents, reaching high recall performance and low serving
latency at the same time. The proposed MEVI provides a unified solution for document retrieval, incorporating the advantages of both embedding-based methods and generation-based methods.
- In our experiments, MEVI significantly improves the recall performance compared to existing methods, achieving  $+3.62\% / +7.32\% / +10.54\%$  of MRR@10/R@50/R@1000 on the MS-MARCO Passage dataset and  $+5.04\% / +5.46\% / +5.96\%$  of R@5/R@20/R@100 on the Natural Questions dataset with low latency capable of serving. Importantly, We also validate MEVI's capability in addressing dynamic corpora effectively.
- We propose a novel RQ-based sequence-to-sequence model to support training and serving on large-scale retrieval datasets. As verified by experiments and analysis, the RQ structure allows for dynamic updates of documents and enables better model performance with low serving latency.
Document retrieval aims to find relevant documents that match a user's query. The core idea of document retrieval is to build an index for a vast number of documents to guarantee fast query response. Document retrieval methods can be divided into three categories.
Traditional sparse retrieval methods build on top of an inverted index using term matching metrics such as TF-IDF [41], query likelihood [23], and a hard-to-beat baseline BM25 [42] in industry-scale web search. Unfortunately, these methods do not take into account the semantics of documents, which can lead to the overlooking of target documents with different wordings.
Dense retrieval methods represent queries and documents using dense embedding vectors, and build ANN index to speed up the search. These methods take advantage of recent advances in pre-trained language models to encode queries and documents into dense representations, such as BERT [8] and RoBERTa [27]. During training, dense retrieval generally follows the contrastive learning paradigm. During inference, ANN methods are applied to search relevant documents, such as tree [1, 25], locality sensitive hashing [6], neighbor graph index (e.g., HNSW [30], DiskANN [43], HMANN [40]), the combination of graph index and inverted index (e.g., SPANN [3]). Sparse retrieval and dense retrieval can be combined [29] to achieve better search quality. However, the discrepancy between KNN and ANN presents a challenge for these methods, particularly in scenarios where there is a significant distribution difference between documents and queries. Besides document retrieval, embeddings are also utilized in cross-modal retrieval [55, 18], graph [5, 19], recommendation [32, 33], etc.
Autoregressive retrieval is a new paradigm to retrieve relevant documents. It uses an end-to-end autoregressive model that takes queries as input and directly generates document identifiers as output. Differentiable Search Index (DSI) [44] and Neural Corpus Indexer (NCI) [46] are the first attempts to propose differentiable indexes for semantics search that directly map queries to relevant document identifiers. DSI generates document identifiers with document tokens, and designs multi-task training for augmentation. NCI builds a tree index by hierarchical KMeans clustering, and represents each file by the path from root to leaf on the tree. It designs a weight-adaptive decoder and uses augmented data to further improve performance. SEAL [2] leverages autoregressive models in another way, generating n-grams and retrieving documents through a pre-built FM-index for n-grams. However, they are all limited by unacceptable serving latency and the inability to handle document updates.
3 Model-enhanced vector index
In this paper, we propose a novel Model-enhanced Vector Index (MEVI) to address the above limitations of both embedding-based dense retrieval methods and generation-based autoregressive retrieval methods. The goal of MEVI is to achieve both fast retrieval speed, like dense retrieval methods, and high recall performance, like autoregressive retrieval methods, by leveraging the strengths of both approaches.
MEVI adopts an autoregressive sequence-to-sequence model as the enhancing model and a well-adapted twin-tower representation model as the vector index. The twin-tower model encodes queries and documents into embedding vectors for common embedding-based retrieval. The autoregressive sequence-to-sequence model takes the same queries as input and outputs the most relevant document clusters. To organize documents into hierarchical clusters, MEVI employs Residual Quantization [31], which first clusters the document embeddings and then iteratively clusters the residual embeddings. During each clustering pass, the residual embeddings are computed by subtracting the cluster centroid embeddings from the embeddings used in the previous pass. RQ is essentially a hierarchical
Figure 1: Overview of Model-enhanced Vector Index (MEVI). (a) Training. Document embeddings and residual embeddings are iteratively clustered to form the RQ structure; then the RQ structure provides cluster codes as labels for the training of sequence-to-sequence model. (b) Inference. The query is fed into the representation model and the sequence-to-sequence model respectively to get candidates from ANN and clusters; then ensemble strategies are performed to get final candidates.
KMeans [15] algorithm, where similar documents are assigned to the same or close clusters. Document clustering not only works in traditional clustering-based retrieval [45, 26], but is also suitable for autoregressive sequence-to-sequence models as mentioned in DSI [44] and NCI [46]. Once the results are obtained from the autoregressive model and the twin-tower model, MEVI ensembles the results to form the final document candidates. The whole workflow is shown in Figure 1.
3.1 Organizing documents with Residual Quantization
Twin-tower representation models [38, 11, 51] are designed for dense retrieval, also known as embedding-based retrieval. They separately encode queries and documents into embeddings, and perform similarity search (e.g. Euclidean, cosine, or inner product) on the embeddings to retrieve relevant documents for given queries. Our method utilizes information from autoregressive sequence-to-sequence models to significantly and accurately narrow down the range of candidate sets. We use the state-of-the-art twin-tower models AR2 [54] and T5-ANCE [51] in our experiments. It is worth noting that the improvement can theoretically be applied to any twin-tower representation models.
To build an autoregressive sequence-to-sequence task, documents need to be organized into some kind of data structures. In DSI [44] and NSI [46], documents are hierarchically clustered through the KMeans algorithm [15] to form a tree, and the T5 model [39] is used to encode the query and decode into a document identifier represented by a path on the tree. The tree structure, however, cannot support large-scale corpus in real applications due to its long sequential decoder steps.
Vector quantization is a useful method to compress document embedding representations. Instead of the commonly used Product Quantization which splits embeddings along dimension, we leverage Residual Quantiztaion (RQ) to preserve the hierarchical semantics of documents. The construction of RQ codebook is a process of iterative clustering. Concretely, assume the initial document embedding is  $X \in \mathbb{R}^{n \times d}$ , the first KMeans clustering is performed on  $X$  to get the centroid embeddings  $B_1 \in \mathbb{R}^{b \times d}$ , also known as codewords, and the one-hot clustering assignment  $C_1 \in \{0, 1\}^{n \times b}$ , also known as codes. Then the second clustering is performed on the residual embeddings  $R_1 = X - C_1B_1$ ,
rather than the original document embeddings  $X$ . Codewords  $B_{2}$  and codes  $C_2$  are obtained in this step. Then the residual  $R_{2} = R_{1} - C_{2}B_{2}$  is fed into the third clustering, and the process continues till a pre-defined number of steps  $m$  is reached. After all the clustering process, the overall codebook is obtained as  $B = [B_{1},B_{2},\dots,B_{m}]\in \mathbb{R}^{mb\times d}$ , and the overall codes for documents are  $C = [C_1^T,C_2^T,\dots,C_m^T ]^T\in \{0,1\}^{n\times mb}$ ; the iterative clustering can be seen as a solution to the minimization problem  $\min ||X - CB||_2^2$ . The explanation of symbols are detailed in Table 1.
Table 1: Notations.
Notation	Explanation
n	The number of documents
d	Embedding dimension
log2b	The bit number of compact codes
b	The number of centroid embeddings for each clustering
m	The number of clustering layers, i.e. the length of cluster-id
X∈Rn×d	All document embedding
Bi∈Rb×d,i=1,...,m	Centroid embeddings of clusters
B∈Rmb×d	Overall centroid embeddings
Ci∈{0,1}n×b,i=1,...,m	Code-style assignment of documents
C∈{0,1}n×mb	Overall cluster assignment
C̅∈{0,1,2,...,b-1}n×m	Overall compact cluster assignment
k	The number of retrieved clusters
K	The number of retrieved documents (for R@K, MRR@K)
ci,i=1,...,k	Top-k retrieved clusters
Ri∈Rn×d,i=1,...,m-1	Residual embeddings after clustering
s0,sc	Scoring function for embedding similarity and cluster ranking
rc	Cluster ranking function
α,β	Hyperparameters of the ensemble process
Both RQ codebooks and hierarchical trees are constructed using iterative clustering. The difference is that RQ clusters the residuals and maintains a fixed step length. Considering that the clustering process is essentially the classification of documents, RQ's clustering of residuals can be regarded as classifying documents in a more and more refined manner. On the other hand, all codes in RQ are of length  $m$ , making the latency of the autoregressive process stable and acceptable.
After the RQ construction process, we obtain the codes of documents  $C$ , or a compact version  $\widetilde{C} \in \{0,1,2,\dots,b - 1\}^{n\times m}$  instead of one-hot vectors, with each code an index of the codebook in the corresponding step. In the decoding process of autoregressive sequence-to-sequence models, the code  $\widetilde{C}$  is generated step by step during inference until step  $m$ , with the vocabulary of each step  $\{0,1,2,\dots,b - 1\}$ . We consider documents to be from the same cluster only if their codes are identical. Therefore, we regard the results from the sequence-to-sequence model as clusters with top-k probability of containing relevant documents. With the mapping of document to code  $\widetilde{C}$ , the sequence-to-sequence model can be trained with a large collection of  $< query, code >$  pairs derived from  $< query, document >$  pairs.
3.2 Ensemble process
Given the top-k retrieved clusters  $\{c_1, c_2, \dots, c_k\}$  from the autoregressive sequence-to-sequence model, we incorporate such information into embedding-based retrieval. The rest of this section explores some possible options to ensemble the top-k cluster information.
Search only in candidates from top-k clusters. A straight-forward strategy is to only search relevant embeddings within top-k clusters. This strategy assumes that the sequence-to-sequence model has already retrieved relevant documents by retrieving top-k related clusters, greatly reducing the number of candidates for embedding similarity search.
Ensemble the original embedding search results with the candidates from the top-k clusters. Considering that the RQ structure may be inaccurate, and the beam search process is a lossy greedy algorithm, the retrieved top-k clusters may not contain all the relevant documents. A more robust
method is to combine candidates from common embedding-based retrieval and top-k clusters. The original similarity score is not applicable for re-ranking these candidates, since the first candidate set already contains the best documents under this similarity metric. Therefore, we design a new score for re-ranking, also leveraging the information of retrieved top-k clusters. Assume  $s_0$  denotes the original scoring function for embedding similarity,  $s_c$  denotes the cluster scoring function of documents,  $\alpha$  denotes the coefficient of cluster scores,  $x$  is a random document, the new score is:
$$
s (x) = s _ {0} (x) + \alpha \cdot s _ {c} (x)
$$
The cluster scoring function can be a function of the cluster ranking:
$$
s _ {c} (x) = \frac {1}{\beta \cdot r _ {c} (x) + 1}
$$
where  $r_c$  denotes the cluster ranking function of documents,  $\beta$  is another hyperparameter. For documents that are not in the top-k clusters, we simply use zero or a value smaller than the minimum cluster score.
In this section, we conduct various experiments to verify the effectiveness and the efficiency of MEVI. We first introduce the datasets and the evaluation metrics to be used (Section 4.1), and then describe the implementation details of MEVI and other baseline methods (Section 4.2). We demonstrate and analyze the effectiveness of MEVI on two large-scale datasets (Section 4.3, 4.4), examine the ability of MEVI to handle unseen documents (Section 4.5), and explore the effectiveness of RQ (Section 4.6). We also examine the efficiency of MEVI (Section 4.7), study the impact of hyperparameters (Section 4.8) and apply our method to practical applications (Appendix A).
4.1 Datasets & metrics
Datasets. We conduct experiments on two widely used large-scale document retrieval benchmarks, i.e. MSMARCO [34] Passage Retrieval dataset and Natural Questions (NQ) [22] dataset. MSMARCO is a question answering dataset formulated by Microsoft in 2016. It consists of real Bing questions and human-generated answers. The version we use is the MSMARCO Passage Retrieval dataset, which contains 8.8 million passages and around 510,000 queries. NQ is a large open-domain question answering dataset containing queries collected from Google search logs. DPR [21] selected around 62,000 factoid questions and processed all the Wikipedia articles (21 million) as the collection of passages. We use the DPR version for evaluation. For both datasets, we use their predetermined training and validation splits for evaluation.
Metrics. To reflect whether relevant documents are retrieved given a user query, we employ Recall@K and MRR@K, two widely accepted information retrieval evaluation metrics. R@K (short for Recall@K) measures the proportion of relevant items found among the retrieved top-K candidates. A higher recall means that there are more relevant documents among the top-K candidates. MRR@K (mean reciprocal rank) calculates the reciprocal of the rank at which the first relevant document is retrieved within top-K candidates. A higher MRR indicates that the relevant document is retrieved with a higher rank.
RQ codebook. By default, we use an RQ codebook with 4 clustering layers and 5 bits per layer. In each layer, we use the default KMeans algorithm provided by scikit-learn [37], where the number of centroids is  $b = 2^5 = 32$ . The entire RQ has  $32^4 = 1,048,576$  clusters, and each cluster contains an average of 8 and 20 documents in two datasets. We tested other configurations of RQ in Section 4.6.
Sequence-to-sequence model. We leverage Neural Corpus Indexer (NCI) [46], the state-of-the-art sequence-to-sequence model for document retrieval. We replace its original decoder vocabularies with our RQ structure, and train NCI upon the processed  $< query, code >$  pairs. For data augmentation, we employ the "document as query" technique and use DocT5Query [35] as suggested in NCI, generating 10 queries per document for MSMARCO Passage, and 1 query per document for NQ. We use the same optimizer and learning rate as NCI; however, we use a larger batch size, 256 per GPU, and disable Rdrop to reduce the overall training cost. For inference, we use beam search with a beam
size of 10 to 1000. All experiments of NCI are conducted on an NVIDIA V100 GPU cluster with 32GB memory per GPU, and each job runs on 8 GPUs.