A Neural Corpus Indexer for Document Retrieval

Introduction

Document retrieval and ranking are two key stages for a standard web search engine [56; 34]. First, the document retrieval stage retrieves candidate documents relevant to the query, and then, the ranking stage gives a more precise ranking score for each document. The ranking stage is often fulfilled by a deep neural network, taking each pair of query and document as input and predicting their relevance score. Nevertheless, a precise ranking model is very costly, while typically only a hundred or thousand candidates per query are affordable in an online system. As a result, the recall performance of the document retrieval stage is very crucial to the effectiveness of web search engines.

Existing document retrieval methods can be divided into two categories, namely term-based and semantic-based approaches . Term-based retrieval approaches [9; 59] build an inverted index for the entire web corpus, but they hardly capture document semantics and fail to retrieve similar documents in different wordings. Thus, semantic-based approaches [56; 36] are proposed to alleviate this discrepancy. First, they learn dense representations for both queries and documents through a twin-tower architecture; then Approximate Nearest Neighbor (ANN) search is applied to retrieve relevant documents for the designated query. Despite of their success in real applications, these approaches can not fully leverage the power of deep neural networks for the following reasons. First, a single embedding vector has limited capacity to memorize all semantics in a document, and it performs even worse than term-based methods in the applications that heavily rely on exact match . Second, the model is unable to incorporate deep query-document interactions. Because ANN algorithms theoretically require a strong assumption for the Euclidean space, we have to adopt simple functions such as cosine similarity to capture the query-document interactions .

Given the above limitations, several research works have explored end-to-end models that directly retrieve relevant candidates without using an explicit index. Gao et al. proposed a Deep Retrieval (DR) framework for item recommendation, which learned a retrievable structure with historical user-item interactions. Nevertheless, it is more challenging to design a universal model for semantic text retrieval, as we need to leverage the power of both pre-trained language models and deep retrieval networks simultaneously. Tay et al. proposed Differentiable Search Index (DSI), a text-to-text model that maps queries directly to relevant docids. To the best of our knowledge, this is the first attempt to propose a differentiable index for semantic search. However, the vanilla transformer decoder in DSI does not fully leverage the hierarchical structures of document identifiers, and the model is pruned to over-fitting with limited training data. Furthermore, Bevilacqua et al. proposed SEAL by leveraging all n-grams in a passage as its identifiers. But for long documents, it is hard to enumerate all possible n-grams. In general, the recall performance of end-to-end document retrieval remains a large room to be improved.

In this paper, we show that the traditional text retrieval frameworks can be fundamentally changed by a unified deep neural network with tailored designs. To this end, we propose a Neural Corpus Indexer (NCI), which supports end-to-end document retrieval by a sequence-to-sequence neural network. The model takes a user query as input, generates the query embedding through the encoder, and outputs the identifiers of relevant documents using the decoder. It can be trained by both ground-truth and augmented query-document pairs. During inference, the top $N$ documents are retrieved via beam search based on the decoder. Designing and training such a model is non-trivial, so we propose several crucial techniques to ensure its effectiveness. First, to get sufficient query-document pairs for training, we leverage a query generation network to obtain possible pairs of queries and documents. Second, we utilize the hierarchical $k$-means algorithm to generate a semantic identifier for each document. Third, we design a prefix-aware weight-adaptive decoder to replace the vanilla one in a sequence-to-sequence architecture. Specifically, the same token will be assigned different embedding vectors at different positions in the identifiers, while another transformer-based adaptive module is applied to the classification weights for token prediction in the context of a certain prefix. This makes the classifiers customized to different prefixes when decoding along the hierarchical tree structure. Besides, a consistency-based regularization loss is taken for training both encoder and decoder networks to mitigate the over-fitting problem.

Our NCI design solves the limitations of traditional index-retrieve pipelines from multiple perspectives. On one hand, a whole neural network model replaces the traditional inverted index or vector search solutions. It can be optimized end-to-end using realistic query-document pairs, which fully captures both term-based and semantic-based features and is adaptive to the changing of workloads. On the other hand, the model is able to capture deep interactions between queries and documents via the encoder-decoder attention, which enlarges the capacity of vector-based representations. Moreover, NCI achieves much better ranking results than ANN-based approaches as it is optimized directly by the final target. Thus, it can be served as an end-to-end retrieval solution while releasing the burden of re-ranking for a long candidate list.

In addition to the superior performance, the invention of Neural Corpus Indexer is also promising from the perspective of system design. As nowadays, ranking and query-answering modules are already implemented by neural networks, NCI finishes the last piece of puzzle for the next-generation information retrieval system based on a unified differentiable model architecture. This reduces the dependency among different sub-modules, while the processes of system deployment and maintenance could be greatly eased.

Our contributions are highlighted as follows.

For the first time, we demonstrate that an end-to-end differentiable document retrieval model can significantly outperform both inverted index and dense retrieval solutions. This finding will inspire research on further steps towards the next-generation search systems, for instance, unifying informational retrieval, ranking, and question answering in a single differentiable framework.

We design a sequence-to-sequence model, named Neural Corpus Indexer (NCI), which generates relevant document identifiers directly for a specific query. In our experiments, the proposed NCI model improves the state-of-the-art performance of existing methods by a significant margin, achieving +21.4% and +16.8% relative enhancement for Recall@1 on NQ320$k$ dataset and R-Precision on TriviaQA dataset, respectively. Also, NCI itself achieves a competitive MRR score without using an explicit ranking model.

We propose a novel decoder architecture, namely prefix-aware weight-adaptive (PAWA) decoder, to generate document identifiers. As verified by ablation studies, this invention is very crucial for NCI to achieve an outstanding performance. Moreover, query generation, semantic document identifiers, and consistency-based regularization are all accountable for the superior capability of Neural Corpus Indexer.

Related work

In this section, we briefly introduce the related works and leave more discussions in Appendix A.

Sparse retrieval. Traditional document retrieval methods are based on Sparse Retrieval, which is built upon inverted index with term matching metrics such as TF-IDF , query likelihood or BM25 . In industry-scale web search, BM25 is a difficult-to-beat baseline owing to its outstanding trade-off between accuracy and efficiency. In recent years, there are some attempts to incorporate the power of neural networks into inverted index. The Standalone Neural Ranking Model (SNRM) learns high-dimensional sparse representations for query and documents, which enables the construction of inverted index for efficient document retrieval. Doc2Query predicts relevant queries to augment the content of each document before building the BM25 index, and DocT5Query improves the performance of query generation by the pre-trained language model T5 . Furthermore, DeepCT calculates context-aware term importance through neural networks to improve the term matching metrics of BM25.

Dense retrieval. Another line of research lies in Dense Retrieval, which presents query and documents in dense vectors and models their similarities with inner product or cosine similarity. These methods benefit from recent progresses of pre-trained language models, such as BERT and RoBERTa to obtain dense representations for queries and documents. At inference time, efficient Approximate Nearest Neighbor (ANN) search algorithms, such as k-dimensional trees , locality-sensitive hashing , and graph-based indexes (e.g., HNSW , DiskANN and SPANN ) can be utilized to retrieve relevant documents within a sublinear time. Besides, Luan et al. analyze the limited capacity of dual encoders, and propose a combination of sparse and dense retrieval methods with multi-vector encoding to achieve better search quality.

Autoregressive retrieval. The other way to approach retrieval is utilizing an end-to-end autoregressive model. Firstly, several efforts have been done on entity linking [13; 12; 11], which can be regarded as a special type of retrieval task, e.g., using an entity to ask the posed question. Recently, different from the entity linking task, Tay et al. proposed the DSI (differentiable search index) model to generate relevant document identifiers directly corresponding to the query. Bevilacqua et al. employed the autoregressive model to generate relevant words for a query and utilize the generated string to retrieve relevant documents. Besides, the Deep Retrieval (DR) approach for recommendation is also related to this category, which learns a deep retrievable network with user-item clicks and gets rid of the ANN algorithms based on the Euclidean space assumption.

Pre-trained language models. Recently, pre-trained Language Models (LMs), such as BERT and RoBERTa , have led to a revolution in web search techniques. The representation vectors for all documents can be calculated and indexed offline. In the online serving stage, it calculates the representation vector for the input query, and applies a crossing layer to calculate the relevance score between each query and document pair. The crossing layer usually adopts simple operators such as cosine similarity or a single feed-forward layer to retain a high efficiency. Gao et al. found that a standard LMs’ internal attention structure is not ready-to-use for dense encoders and proposed the Condenser to improve the performance of dense retrieval. Moreover, ANCE leverages hard negatives to improve the effectiveness of contrastive learning, which generates better text representations for the retrieval tasks.

Neural corpus indexer

The neural corpus indexer (NCI) is a sequence-to-sequence neural network model. The model takes a query as input and outputs the most relevant document identifier (docid), which can be trained by a large collection of pairs. The documents are encoded into semantic docids by the hierarchical $k$-means algorithm , which makes similar documents have “close” identifiers in the hierarchical tree. As shown in Figure 1, NCI is composed of three components, including Query Generation, Encoder, and Prefix-Aware Weight-Adaptive (PAWA) Decoder. Query generation is implemented by a sequence-to-sequence transformer model that takes as input the document terms and produces a query as output . The encoder, following the standard transformer architecture, is composed of $M_{1}$ stacked transformer blocks, which outputs the representation for an input query. For the decoder network, we stack $M_{2}$ transformer layers. To better align with the hierarchical nature of the semantic identifiers, we propose a weight adaptation mechanism based on another transformer to make the decoder aware of semantic prefixes. At inference time, the top $N$ relevant documents can be easily obtained via beam search. Due to the hierarchical property of semantic identifiers, it is easy to constrain the beam search on the prefix tree so that only valid identifiers will be generated.

NCI generates document identifiers solely based on the input query without explicit document content, which is difficult when the size of the corpus is very large. Thus, we aim to inject useful priors into the identifiers so that the semantic information of documents can be incorporated in the decoding process. In other words, we hope the documents with similar semantics have close docids to facilitate the learning process of NCI. To achieve this, we leverage the hierarchical $k$-means algorithm to encode documents. As shown in Figure 1(a), given a collection of documents to be indexed, all documents are first classified into $k$ clusters by using their representations encoded by BERT . For cluster with more than $c$ documents, the $k$-means algorithm is applied recursively. For each cluster containing $c$ documents or less, each document is assigned a number starting from 0 to at most $c$-1. In this way, we organize all documents into a tree structure $T$ with root $r_{0}$. Each document is associated with one leaf node with a deterministic routing path $l=\{r_{0},r_{1},...,r_{m}\}$ from the root, where $r_{i}\in[0,k)$ represents the internal cluster index for level $i$, and $r_{m}\in[0,c)$ is the leaf node. The semantic identifier for a document is concatenated by the node indices along the path from root to its corresponding leaf node. For documents with similar semantics, the prefixes of their corresponding identifiers are likely to be the same. For simplicity, we set $k=30$ and $c=30$ in all experiments, leaving the optimization of these hyper-parameters to future work. The detailed procedure of hierarchical $k$-means will be described in Algorithm 1 in the Appendix B.2.

2 Query generation

One challenge of generating document identifiers by single query input is how to make the identifiers aware of the document semantics. Since the content of each document is not explicitly known at inference, it must be incorporated into the model parameters during training. To facilitate the training process, we generate a bunch of queries with a query generation module and bind the information of document content through training the sequence-to-sequence model with generated queries and their corresponding document identifiers. In NCI, we utilize two kinds of augmented queries:

DocT5Query. We adopt a standard sequence-to-sequence transformer based on the implementation of DocT5Query pre-trained by a large query-document corpus. It takes as input the document terms and produces relevant queries via random sampling. Note that we use random sampling instead of beam search to ensure the diversity of generated queries.

Document As Query. Like DSI , we also utilize the first 64 terms for each document as queries. Besides, we randomly selected 10 groups of 64 consecutive terms from the whole article as additional queries. This makes the NCI model aware of the semantic meaning of each document.

3 Prefix-aware weight-adaptive decoder

Given an input query $x$, the probability of generating a document identifier can be written as:

where $r_{i}$ is the $i$-th token in the current identifier; $x$ is the representation output from encoder; $\theta$ denotes the total parameters and $\theta_{i}$ is the parameter for the $i$-th step.

This probability can be modeled by a transformer-based decoder. For an internal node with level $i$, the probability is calculated by:

As the encoder and decoder utilize distinct vocabulary spaces, we do not share the embedding space for their tokens. Different from a standard decoding task, the meanings of the same token appearing at different places of the same identifier are different, as they correspond to different clusters in the hierarchical tree structure. For instance, the “$5_{\underline{2}}$” and “$5_{\underline{3}}$” of the same identifier “$3_{\underline{1}}5_{\underline{2}}5_{\underline{3}}$” correspond to different semantic meanings. Moreover, the same token in the same position may have different semantics with different prefixes. For example, in identifiers “$1_{\underline{1}}1_{\underline{2}}5_{\underline{3}}$” and ”$2_{\underline{1}}4_{\underline{2}}5_{\underline{3}}$”, the same token “$5_{\underline{3}}$” has different semantics in two different identifiers, as they are routed from different prefix paths. These two properties of the hierarchical semantic identifiers motivate us to design the novel Prefix-Aware Weight-Adaptor (PAWA) decoder.

Unlike a standard transformer decoder, the probabilities at different tree levels, such as $p(r_{i}|x,r_{1..{i-1}},\theta_{i})$ and $p(r_{j}|x,r_{1..{j-1}},\theta_{j})$ where $i\neq j$, do not share parameters with each other. To distinguish different semantic levels, we concatenate the position and token values as input for each decoding step, as shown in the left corner of Figure 2. Specifically, we have “$(1,3)(2,5)(3,5)$” for the semantic identifier “$3_{\underline{1}}5_{\underline{2}}5_{\underline{3}}$”, while “$(2,5)$” and “$(3,5)$” represent different tokens in the vocabulary space. As the token embedding and linear classification layers share the same weights, the same token value in different positions would correspond to different model parameters. Moreover, to reflect the influence of different prefixes, we expect the linear classification layer to be aware of different prefixes for predicting a specific token. Concretely, instead of using the same projection weight $W$ in the linear classification layer, we employ the prefix-aware adaptive weights for each token classifier, which can be calculated by another transformer decoder,

For instance, to predict the third tokens in the identifiers “(1,3)(2,1)(3,5)” and “(1,2)(2,4)(3,5)”, respectively, the corresponding adaptive weights are derived separately for different prefixes, i.e., “(1,3)(2,1)” and “(1,2)(2,4)”. As we already know the previous tokens for each position in the teacher forcing setting, the prefix-aware adaptive weights can be calculated and trained in parallel in different positions while adding little burden to the entire model.

4 Training and inference

Consistency-based regularization. To alleviate over-fitting, we employ a consistency-based regularization loss for training each decoding step. Given an input query $q$, we denote the decoder representations by two forward passes with independent dropouts before Softmax as $\mathbf{z}_{i,1}=D(r_{i}|E(q),r_{1,...,i-1},\theta_{i})$ and $\mathbf{z}_{i,2}=D(r_{i}|E(q),r_{1,...,i-1},\theta_{i})$, respectively, where $E(\cdot)$ denotes the encoder network and $D(\cdot)$ denotes the decoder network. The consistency-based regularization loss tries to distinguish the representations from the same token from those of other tokens, like contrastive learning . The regularization loss of query $q$ for the $i$-th decoding step is defined as,

where we leverage dot-product for $sim(\cdot)$; $Q$ is the number of queries in the batch, and the temperature parameter is set as $\tau=1$ in all the experiments.

Training loss. Given a set of training examples $\mathcal{D}=\{(q,d)\}$ composed of queries (training queries and augmented queries) and document identifiers, the loss function can be written as follows:

where $p(d|E(q),\theta)$ denotes the probability of generating $d$ with $q$ as the input. The first part is the seq2seq cross-entropy loss with teacher forcing and the second part is the consistency-based regularization loss summed by all decoding steps. The whole process formulates a sequence-to-sequence neural network, which can be optimized end-to-end via gradient descent. The hyper-parameter $\alpha$ denotes a scaling factor of regularization loss, which will be analyzed in Section 4.4.

Inference via beam search. In the inference stage, we calculate the query embedding through the encoder network and then perform beam search on the decoder network. Due to the hierarchical nature of docid, it is convincing to constrain the beam search decoding process with a prefix tree, which in turn only generates the valid identifiers. The time complexity of beam search is $O(LBF)$, where $L$ is the max length of identifiers (the depth of tree), $B$ is the beam size and $F$ is the max fanout of the tree (30 in our experiments). Given a balanced tree structure built by a corpus with $N$ documents, the average time complexity for beam search is $O(B\text{log}N)$. We leave detailed descriptions of the constrained beam search algorithm in Appendix B.3.

Experiments

In this section, we empirically verify the performance of NCI and the effectiveness of each component on the document retrieval task, which generates a ranking list of documents in response to a query. In the following, we discuss the datasets and evaluation protocols in Section 4.1, describe the implementation details and baseline methods in Section 4.2, and present empirical results and analyses in Section 4.3 and 4.4, respectively.

Datasets. We conduct our experiments on two popular benchmarks for document retrieval, i.e., the Natural Questions and TriviaQA dataset . Natural Questions (NQ) was introduced by Google in 2019. The version we use is often referred to as NQ320$k$, which consists of 320$k$ query-document pairs, where the documents are gathered from Wikipedia pages and the queries are natural language questions. We use its predetermined training and validation split for evaluation. TriviaQA is a reading comprehension dataset , which includes 78$k$ query-document pairs from the Wikipedia domain. Unlike the NQ320$k$ dataset, a query may include multiple answers in TriviaQA.

Metrics. We use widely accepted metrics for information retrieval, including Recall@$N$, Mean Reciprocal Rank (MRR) and R-precision. Recall@$N$ measures how often the desired document is hit by the top-$N$ retrieved candidates. MRR calculates the reciprocal of the rank at which the first relevant document is retrieved. R-Precision is the precision after $R$ documents have been retrieved, where $R$ is the number of relevant documents for the query. A high recall means that the ground truth document is contained in the retrieved candidate list, while a high MRR indicates that the corresponding document has already been ranked at the top position without re-ranking.

2 Implementation details

Hierarchical semantic identifier. For semantic identifiers, we apply a hierarchical $k$-means algorithm over the document embeddings obtained through a 12-layers BERT model with pre-trained parameters (provided by HuggingFace ). For each hierarchical layer, we employ the default $k$-means algorithm implemented in scikit-learn with $k=30$. For simplicity, the recursion terminal condition is also set as $c=30$.

Query generation. We leverage the pre-trained model, DocT5Query , for query generation. We provide all document contents in NQ320$k$ and TriviaQA datasets to predict augmented query-document pairs. For each document, we generate $15$ queries with the first $512$ tokens of the document as input and constrain the maximum length of the generated query as 64.

Training and inference. The Neural Corpus Indexer is implemented with python 3.6.10, PyTorch 1.8.1 and HuggingFace transformers 3.4.0. We utilize the parameters of the T5 pre-trained model to initialize the encoder and randomly initialize the PAWA decoder. All NCI experiments are based on a learning rate $2\times 10^{-4}$ for the encoder and $1\times 10^{-4}$ for the decoder with a batch size $16$ per GPU. We set the scaling factor of the consistency-based regularization loss as $\alpha=0.15$ and the dropout ratio as $0.1$. For inference, we apply the partial beam search algorithm to the trained seq2seq model. We set the length penalty and the beam size as $0.8$ and $100$, respectively. All experiments are based on a cluster of NVIDIA V100 GPUs with 32GB memory. Each job takes 8 GPUs, resulting in a total batch size of $128$ ($16\times 8$).

Baselines. We evaluate BM25 on both raw documents and those augmented by DocT5Query by an open-source implementation . The performance of DSI is referred from its original paper as the implementation has not been officially open-sourced. To avoid the difference in data processing, we reproduce SEAL and ANCE by their official implementations. Some baselines for the TriviaQA dataset are directly referred from . We leave the detailed settings in Appendix B.4.

3 Results

In Table 1 and 2, we compare the empirical results of NCI and corresponding baselines on two benchmarks. We report NCI models based on T5-Base, T5-Large, and ensemble architectures. One can see that even with the T5-Base architecture, NCI outperforms all baselines by a significant margin across four different metrics on both the NQ320$k$ and TriviaQA datasets. Furthermore, an ensemble of five NCI models also brings a large enhancement, because each model is trained individually with a separate semantic identifier generated by a random k-means initialization, making the models complementary to each other. Expect for NCI, SEAL achieves the second best performance. This verifies the superiority of deep text retrieval over traditional sparse and dense retrieval methods. Comparing to SEAL, NCI improves $17.6\%$ for Recall@1, $10.0\%$ for Recall@10, $3.2\%$ for Recall@100, and $14.9\%$ for MRR@100 on the NQ320$k$ dataset. We find that the generated queries have different distributions with the training queries , so we also fine-tune Doc2Query on this dataset for a comparison (denoted by w/ qg-ft). Finally, we achieve 72.78% for Recall@1, outperforming SEAL by 21.4%. On the TriviaQA dataset, NCI obtains $7.9\%$ improvement for Recall@5, $5.5\%$ for Recall@20, $6.0\%$ for Recall@100, and $16.8\%$ for R-Precision. As shown in ablation studies, these improvements are owning to the novel designs of PAWA decoder, query generation, semantic identifiers, and consistency-based regularization. We also notice that query generation plays a key role in boosting the retrieval performance. With query generation, the BM25 + DocT5Query method achieves higher performance than the vanilla BM25, especially on the NQ320$k$ dataset. ANCE achieves competitive performance after fine-tuned by the training pairs, but the performance is relatively lower than our NCI model. Moreover, the MRR@100 and R-Precision metrics of NCI are outstanding, indicating that 80% of the queries can be fulfilled without re-ranking on the retrieved document list. This demonstrates the potential of NCI to be served as an end-to-end solution that replaces the entire index-retrieve-rank pipeline in traditional web search engines.

Furthermore, to study the effect of each component, we report ablation results on both NQ320$k$ and TriviaQA datasets in Table 3. In general, all five components are able to improve the performance of document retrieval, which are detailed below.

w/o DocT5Query. This configuration removes the training queries generated by DocT5Query. According to the results, the query generation model greatly boosts the performance. The result is aligned with our expectation because training with augmented queries allows the NCI model to better understand the semantic meanings of each document.

w/o document as query. Similar to DSI , using the document contents as queries also makes the model aware of the semantics of documents.

w/o PAWA decoder. This configuration removes the adaptive decoder layer in Equation (4) and leverages shared weights with token embedding for the linear classification layer. We notice that the prefix-aware weight-adaptive decoder has a noticeable influence on the performance, which indicates that, instead of borrowing the vanilla transformer decoder, it is necessary to design a tailored decoder architecture for the task of semantic identifier generation.

w/o semantic id. This configuration replaces the semantic identifier of each document to a randomly generated one. We find a relative drop in the model performance on all four metrics, demonstrating that the semantic identifiers derived by the hierarchical $k$-means have injected useful priors. We conjecture that the performance enhancement would be more significant on a larger document corpus.

w/o regularization. There is a performance drop on all four metrics without using consistency-based regularization loss. The reason is that the decoder network is prone to over-fitting. By making the prediction results of two augmented queries consistent, the decoder will become more generalizable and resistant to over-fitting.

w/o constrained beam search. This configuration disables the validating constraint in beam search. In other words, the decoder network does not have a tree-based prior structure. Instead, all tokens in the vocabulary can be generated in each decoding step. We observe a performance drop on four evaluation metrics. This indicates that it is difficult to remember all information of valid identifiers in the network, and an explicit prior could be helpful for improving the quality of beam search.

4 Analysis

Model capacity. Figure 3 compares the learning curves of NCI with different model capacities, which are identical to the small, base, and large settings of ordinary T5 . We observe that with the increase of model size, NCI convergences more quickly with fewer epochs. At convergence, the small model achieves a relatively lower recall. Instead, both the base and large models achieve similar results after sufficient training epochs, and the large model will be slightly higher. This implies that the model capacity has a critical impact on the retrieval performance, and the capacity of base model seems to be enough to memorize all documents in NQ320$k$ and TriviaQA datasets. The large model can be used when the computation capacity is sufficient. For a larger corpus, one may need to increase the model size to obtain satisfactory performance.

Layer number of PAWA adapter. We study the influence of the number of transformer layers in the PAWA adapter and choose the layer number from {0,1,2,4,6,8}. The results are summarized in Table 4. We notice that with the increase of layer number, i.e. from 0 to 4, the overall performance is consistently improved on four metrics. But when the number of layers achieves 6, the performance decreases. When continuing to increase the number of layers to 8, the performance drops significantly. We attribute that to the overfitting issue caused by a large PAWA decoder. Therefore, we adopt the PAWA decoder with a 4-layers adapter in NCI.

Retrieved documents and their semantics identifiers. To verify the effectiveness of retrieval as well as the semantic identifiers learned by the hierarchical $k$-means, we analyze the retrieval results of NCI for some exemplar queries. To illustrate, we select four queries denoted by A-1, A-2, B-1 and B-2, where two queries inside the same group are semantically similar, and the queries in different groups correspond to distinct topics. In Figure 4, we show the probabilities of retrieved documents for each query in group A and B, respectively. The digits along x-axis denote the four-bit prefixes for the semantic identifiers of retrieved documents, and the y-axis stands for their probabilities. We notice that similar queries result in close document distributions, while dissimilar queries in different groups result in un-overlapped document collections. In addition, the documents retrieved by the same group of queries have close prefixes for the identifiers, e.g., 6030, 6032, 6033, 6034 in group A and 7511, 7514, 7516 in group B. Also, we visualize the BERT-based document embeddings by t-SNE in Figure 4, in which each color represents the corresponding documents for a specific query. As shown in the figure, these documents naturally form two clusters with respect to different query groups. Thus, we conclude that the semantic document identifiers generated by the hierarchical $k$-means algorithm have positive effects on the retrieval performance.

Efficiency Analysis. We use an NVIDIA V100-32G GPU to analyze the efficiency of NCI. As the inference speed is influenced by both model capacity and beam size, we report the latency and throughput measures for multiple settings in Table 5. As NCI is an end-to-end retrieval method and achieves competitive performance without re-ranking, the latency and throughput are already affordable for some near-real-time applications. The latency of NCI is on par with DSI and SEAL using the same model size and beam size, because all of them conduct beam search based on transformer decoders. BM25 is very efficient (<100ms per query on CPU using an open-source implementation ), but the recall metrics are much lower. Furthermore, we can leverage other techniques to improve the efficiency of NCI, which will be discussed in the later section.

Limitation & Future Works

Despite the significant breakthrough, the current implementation of NCI still suffers from several limitations before deployment in a large-scale search system. Firstly, it requires a much larger model capacity for extending NCI to the web scale. Secondly, the inference speed needs to be improved to serve online queries in real time. Thirdly, it is difficult to update the model-based index when new documents are added to the system. In future works, we may tackle these problems from four aspects. (1) The architecture of sparsely-gated Mixture of Expert (MoE) can be employed to enhance the model capacity. (2) Documents can be grouped into semantic clusters, and NCI can be used to retrieve relevant cluster identifiers. In this way, all documents in relevant clusters can be retrieved efficiently. (3) Model compression techniques, like weight quantization and knowledge distillation , can be further taken to speed up inference. (4) We plan to explore a hybrid solution by building another index that serves new documents through traditional indexing algorithms.

Conclusion

In this work, we introduce a novel document retrieval paradigm that unifies the training and indexing stages by an end-to-end deep neural network. The proposed Neural Corpus Indexer (NCI) directly retrieves the identifiers of relevant documents for an input query, which can be optimized end-to-end using augmented query-document pairs. To optimize the recall and ranking performance, we invent a tailored prefix-aware weight-adaptive decoder. Empirically, we evaluate NCI on NQ320$k$ and TriviaQA datasets, demonstrating its outstanding performance over state-of-the-art solutions.

References

Appendix A Related work

Traditional web search techniques follow a two-stages paradigm including document retrieval and document ranking. The first stage aims to select a collection of documents relevant to a given query, which requires an ingenious trade-off between efficiency and recall. Then, the document ranking stage takes more advanced features and deeper models to calculate a fine-grained ranking score for each query and document pair. In the following, we first discuss related works for document retrieval and ranking respectively. Afterwards, we introduce recent works that incorporate pre-trained language models into these two stages. At last, the attempts on end-to-end retrieval will be discussed.

Traditional document retrieval methods are based on Sparse Retrieval, which is built upon inverted index with term matching metrics such as TF-IDF , query likelihood or BM25 . In industry-scale web search, BM25 is a difficult-to-beat baseline owing to its outstanding trade-off between accuracy and efficiency. In recent years, there are some attempts to incorporate the power of neural networks into inverted index. The Standalone Neural Ranking Model (SNRM) learns high-dimensional sparse representations for queries and documents, which enables the construction of inverted index for efficient document retrieval. Doc2Query predicts relevant queries to augment the content of each document before building the BM25 index, and DocT5Query improves the performance of query generation by the pre-trained language model T5 . Furthermore, DeepCT calculates context-aware term importance through neural networks to improve the term matching metrics of BM25.

Another line of research lies in Dense Retrieval, which presents query and documents in dense vectors and models their similarities with inner product or cosine similarity. These methods benefit from recent progresses of pre-trained language models, such as BERT and RoBERTa to obtain dense representations for queries and documents. At inference time, efficient Approximate Nearest Neighbor (ANN) search algorithms, such as k-dimensional trees , locality-sensitive hashing , and graph-based indexes (e.g., HNSW , DiskANN and SPANN ) can be utilized to retrieve relevant documents within a sublinear time. Besides, Luan et al. analyze the limited capacity of dual encoders, and propose a combination of sparse and dense retrieval methods with multi-vector encoding to achieve better search quality.

A.2 Document ranking

Document ranking has been extensively studied in recent years and experienced a huge improvement with the booming of deep neural networks. Neural network-based document ranking models mainly fall into two categories. Representation-based models like DSSM (Deep Structured Semantic Model) and CDSSM (a convolution-based variant of DSSM) represent query and document in a shared semantic space and model their semantic similarity through a neural network. In contrast, Interaction-based models first build interactions between query and document terms, and then utilizes neural networks to learn hierarchical interaction patterns. For example, DRMM (Deep Relevance Matching Model) extracts interactive features by matching histograms and utilizing a feed forward network with term-gating mechanism to calculate the relevance score of a query-document pair.

A.3 Pre-trained language models

Recently, Pre-trained Language Models (PLMs) like BERT have led to a revolution of web search techniques. The vanilla BERT model utilizes a single-tower architecture that concatenates query and document tokens as a whole input to the relevance model. Despite of its superior performance, the high computational cost hinders its application to industrial-scale web search systems. TwinBERT tackles this problem by exploiting a Siamese architecture, where queries and documents are first modeled by two BERT encoders separately, and then an efficient crossing layer is adopted for relevance calculation. The representation vectors for all documents can be calculated and indexed offline. In the online serving stage, it calculates the representation vector for the input query and applies a crossing layer to calculate the relevance score between each query and document. The crossing layer usually adopts simple similarity functions such as dot product or a single feed-forward layer to achieve a high efficiency.

Moreover, Chang et al. argue that the Masked Language Model (MLM) loss designed for BERT pre-training is not naturally fitted to embedding-based retrieval tasks. Instead, they propose three paragraph-level pre-training tasks, i.e., Inverse Cloze Task (ICT), Body First Selection (BFS), and Wiki Link Prediction (WLP), which demonstrate promising results in text retrieval experiments. Gao et al. find that a standard LMs’ internal attention structure is not ready-to-use for dense encoders. Thus, they propose a novel architecture named Condenser to improve the performance of dense retrieval. ANCE (Approximate nearest neighbor Negative Contrastive Estimation) leverages hard negatives to improve the effectiveness of contrastive learning, which generates better text representations for the retrieval task.

A.4 End-to-end retrieval

The deficiency of index-retrieve paradigm lies in that the two stages of document retrieval and re-ranking are optimized separately. Especially, the document retrieval procedure is often sub-optimal and hinders the performance of the entire system. Thus, there are some recent attempts to achieve end-to-end retrieval as a one-stage solution. ColBERT introduces a contextualized late interaction architecture, which independently encodes query and document through BERT, and performs cross-term interaction based on the contextualized representations of query and document terms. ColBERT supports end-to-end retrieval directly from a large document collection by leveraging vector-similarity indexes in the pruned interaction layer. It can be viewed as a compromise between single-tower and twin-tower BERT architectures which maintains an effective trade-off between accuracy and latency. Moreover, the Contextualized Inverted List (COIL) exacts lexical patterns from exact matching pairs through contextualized language representations. At search time, we build representation vectors for query tokens and perform contextualized exact match to retrieve relevant documents based on inverted index.

Although ColBERT and COIL have shown promising results in end-to-end retrieval tasks without re-ranking, their performance is still not obviously better (if not worse) than a common practice of “BM25 indexer + BERT re-ranker”, and their efficiency is also not good enough for an industrial web search engine. Therefore, we resort to a new indexing paradigm to break the bottleneck. We believe the neural corpus indexer proposed in this paper is a crucial break-through, opening up new opportunities to optimize the performance of web-scale document retrieval. Moreover, there are a few attempts that try to build a model-based search index by directly predicting document identifiers. Tay et al. proposed the DSI (differentiable search index) model based on an encoder-decoder architecture to generate relevant docids. However, its decoder architecture remains the same as T5, which is unsuitable to generate semantic ids derived by hierarchical $k$-means. SEAL uses all n-grams in a passage as its possible identifiers and build a FM-Index to retrieve documents; but it is hard to enumerate all n-grams for retrieving relevant documents. In addition, our work is related to Deep Retrieval for the recommendation task, which learns a deep retrievable network with user-item clicks without resorting to ANN algorithms constrained by the Euclidean space assumption.

Appendix B Reproducibility

We provide our code for reproduction in the supplementary material. We will release it to public shortly.

We conduct experiments on NQ320$k$ and TriviaQA datasets. For NQ320$k$ dataset, the queries are natural language questions and the documents are Wikipedia articles in HTML format. During dataset processing, we first filter out useless HTML tag tokens, and extract title, abstract and content strings of each Wikipedia article using regular expression. The experiments are also conducted on TriviaQA dataset. For TriviaQA dataset, it includes 78$k$ query-document pairs from the Wikipedia domain, which are processed almost the same as NQ320$k$. Then, we detect duplicated articles based on the title of each article. After that, we concatenate the title, abstract and content strings of each Wikipedia article, and apply a 12 layers pre-trained BERT model on it to generate document embeddings. Finally, hierarchical $k$-means is applied on the article embeddings to produce semantic identifiers for each article.

B.2 Hierarchical k𝑘k-means for semantic identifier

The pseudo code of hierarchical $k$-means is detailed in in Algorithm 1.

B.3 Constrained beam search

The pseudo code of constrained beam search is detailed in Algorithm 2.

B.4 Baselines

We describe the baseline methods in this section. For most of them, we use their official open-source implementations.

BM25. BM25 is currently the mainstream algorithm for calculating the similarity score between query and document in information retrieval . We calculate BM25 between an original query $Q$ and a document $d$ which derived from a sum of contributions from each query term $q_{i}$ as,

where $w_{i}$ denotes the weight of $q_{i}$, and $R(q_{i},d)$ is the correlation between $q_{i}$ and $d$. We use the open-source implementation from Rank-BM25 https://github.com/dorianbrown/rank_bm25.

BM25 + DocT5Query. The docT5Query model generates questions that related to a document. These predicted queries are then appended to the original documents, which are then indexed. Note that we use the same predicted queries in our query generation module. Queries are issued against the index as “bag of words” queries, using BM25 for evaluation. We use the open-source code for DocT5Queryhttps://github.com/castorini/docTTTTTquery, and the generated queries keep the same with NCI (our model) to have a fair comparison.

BERT + ANN (Faiss). We use the Flat Index method with the query and document representations obtained by CoCondenserhttps://github.com/luyug/Condenser which is pretrained on Wikipedia and then finetuned over NQ dataset. For the Flat Index method, we use the version implemented by Faisshttps://github.com/facebookresearch/faiss.

BERT + BruteForce. In this baseline, we use the CoCondenser , pretrained on Wikipedia and then finetuned over NQ dataset, to encode queries and documents separately. Then, the Cosine Similarity is computed for each query and document pair. After that, for each query, the documents with the largest Cosine Similarity score are retrieved.

ANCE (MaxP & FirstP). ANCE, a training mechanism, that constructs negatives from an Approximate Nearest Neighbor (ANN) index of the corpus . For BERT FirstP, we concatenate the title and content of each document by a [SEP] token. For BERT MaxP, we only use the content of each document. We use the open-source implementationhttps://github.com/microsoft/ANCE.

SEAL (BART-Large). We reproduce SEAL based on the open-sourced implementationhttps://github.com/facebookresearch/SEAL.

DSI. The DSI model learns a text-to-text model that maps string queries directly to relevant docids . We report the performance of DSI (T5-Base), DSI (T5-Large) and DSI (T5-XXL) from its original paper as the implementation has not been open-sourced.

Appendix C More Experimental Results

We study the influence of regularization strength and choose the regularization hyper-parameter $\alpha$ from {0, 0.1, 0.15, 0.2, 0.3}. Table 6 summaries the results with different regularization hyper-parameter $\alpha$ settings. At convergence, the hyper-parameter $\alpha$ = 0.15 generally achieves better performance. Therefore, we set the default value as $\alpha$ = 0.15 in NCI.

Appendix D Miscellaneous

This work aims at introducing a new learning paradigm that can unify the learning and indexing stages with an end-to-end deep neural network. Besides, our work has the potential to inspire more attempts at unifying the retrieval and re-ranking task with an end-to-end framework, which might have positive social impacts. We do not foresee any form of negative social impact induced by our work.

Privacy Information in Data.

We use the NQ dataset privided by the work . The dataset only includes questions, rendered Wikipedia pages, tokenized representations of each page, and the annotations added by our annotators. No privacy information is included. For the TriviaQA , which is a reading comprehension dataset, it includes 78$k$ query-document pairs from the Wikipedia domain. Again, no privacy information is included.