A Neural Corpus Indexer for Document Retrieval
Authors: Yujing Wang
Yingyan Hou
Haonan Wang
Ziming Miao
Shibin Wu
Hao Sun
Qi Chen
Yuqing Xia
Chengmin Chi
Guoshuai Zhao
Zheng Liu
Xing Xie
Hao Sun
Weiwei Deng
Qi Zhang
Mao Yang
Venue: NeurIPS 2022 Accept (2022)
Keywords: document retrieval
sequence-to-sequence
model-based index

Abstract: Current state-of-the-art document retrieval solutions mainly follow an index-retrieve paradigm, where the index is hard to be directly optimized for the final retrieval target. In this paper, we aim to show that an end-to-end deep neural network unifying training and indexing stages can significantly improve the recall performance of traditional methods. To this end, we propose Neural Corpus Indexer (NCI), a sequence-to-sequence network that generates relevant document identifiers directly for a designated query. To optimize the recall performance of NCI, we invent a prefix-aware weight-adaptive decoder architecture, and leverage tailored techniques including query generation, semantic document identifiers, and consistency-based regularization. Empirical studies demonstrated the superiority of NCI on two commonly used academic benchmarks, achieving +21.4% and +16.8% relative enhancement for Recall@1 on NQ320k dataset and R-Precision on TriviaQA dataset, respectively, compared to the best baseline method.

A Neural Corpus Indexer for Document Retrieval
Yujing Wang'
Yingyan 12-
Haonan Wang Ziming Miao'
Shibin 2u
Hao Sun 1,4x
Qi Chen' Yuqing Xia'
Chengmin Chi Guoshuai Zhao' Zheng Liu
Xing Xie Hao Allen Sun! Weiwei Deng' Zhang' Mao Yang'
Microsoft Tsinghua University University of Illinois, Urbana Champaign *Peking University
fyujwang, zimiao, cheqi, yuqxia, chec, zhengliu) crrosoottttoooooocooo fguzhao, xingx, hasun, dedeng, zhang.qi qi, maoyang) www.iiirrooo..oommm com
2 {Ohyy20, wusb20)@mails tsinghua. edu. cn haonan30lllimois.. edu sunhao@stu, pku. edu.
Abstract
Current state-of-the-art document retrieval solutions mainly follow an index- retrieve paradigm, where the index is hard to be directly optimized for the final retrieval target. In this paper, we aim show that end-to-end deep neural network unifying training and indexing stages can significantly improve the recall performance of traditional methods. To this end, propose Neural Corpus In- network that generates relevant document
dexer (NCI), sequence-lo-seuee identifiers directly for a designated query. To optimize the recall performance of NCI, we invent a prefix-aware weight-adaptive decoder architecture, and leverage tailored techniques including query generation, semantic document identifiers, and consistency-basee regularization. Empirical studies demonstrated the superiority of NCI on two commonly used academic benchmarks, achieving +21.45 1.4% and 1 16.8% relative enhancement for Recall @ 1 on NQ320k dataset and R-Precision on TriviaQA dataset, respectively, compared to the best baseline method.
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 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 affordable in 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 comnantic-based approaches [22]. Term- based retrieval approaches [9, 59] build an inverted index for the entire web corpus, but they hardly capture document semantics and fail 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 twin-tower architecture; then Approximate Nearest Neighbor (ANN) search is applied to retrieve relevant documents for the designated query. Despite of their success 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 document, and it performs even worse than term tered methods in the applications that heavily rely on exact match [37]. Second, the model is unable incorporate deep query-document interactions. Because
TThe work was done Microsoft.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

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 [20].
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. [20] proposed a Deep Retrieval (DR) framework for item recommendation, which learned a retrievable structure with historical user-item interactions. Nevertheless, it 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. [50] proposed Differentiable Search Index (DSI), text-to-texi 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. [4] 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 model takes user query as imput, 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 -trth and augmented query-document pairs. During inference, the top N documents are retrieved via beam search based on the decoder. Designing and training such model non-trivial, so we propose several crucial techniques to ensure its effectiveness. First, to get sufficient query docymed 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 semantic identifier for each document. Third, we design a prefix-aware weight-adaptive decoder to replace the vanilla one in a seuuence-to-sqqunne architecture. Specifically, the same token will be assigned different embedding vectors at different positions in the identifiers, while another gransformer-basee adaptive module is applied to the classification weights for token prediction in the context of certain prefix. This makes the classifiers customized to different prefixes when decoding along the hicrarchical tree structure. Besides, a consistency-haid based regularization loss is taken for training both encoder and decoder networks to mitigate the over-fitting problem.
neural network. The
Our NCI design solves the limitations of traditional index-retrieve pipelines from multiple perspec- tives. 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- term-based and semantic-basse bbasd 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- vector-based representations. Moreover, NCI achieves much better ranking results than ANN-based approaches 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 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 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 diffcrentiable document retrieval model 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 single differentiable framework.
We design a relevant document identifiers directly for a specific query. In our experiments, the proposed NCI model improves the sfate-of-the-ar performance of existing methods by significant margin, achieving +21.4% and +16.8% relative enhancement for Recall @ on NQ320k dataset and
model, named Neural Corpus Indexer (NCI), which generates
2

---

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-adaplive (PAWA) decoder, to generate document identifiers. As verified by ablation studies, this invention is very crucial for NCI to achieve outstanding performance. Moreover, query generation, semantic document identifiers, and consistency- regularization are all accountable for the superior capability of Neural Corpus Indexer.
2 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 metries such as TF-IDF [45]. query likelihood [33] or BM25 [44]. In industry-scale web search, BM25 is a difficult-to-beal 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) [57] learns high high-dimensional sparse representations for query and documents, which enables the construction of inverted index for efficient document retrieval. Doc2Query [41] predicts relevant queries augment the content of each document before building the BM25 index, and DocTSQuery [40] improves the performance of query generation by the pre-trained language model T5 [5]. Furthermore, DeepCT [9] 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- trairid language models, such BERT [14] and RoBERTa [35] to obtain dense representations for queries and documents. At inference time, efficient Approximate Nearest Neighbor (ANN) search algorithms, such as k-dimensional trees [3]. locality-sensitive
hashing [10], and graph-based indexes (e.g., HNSW [38], DisKANN [27] and SPANN (7]) can be
utilized to retrieve relevant documents within a sublinear time. Besides, Luan et al. [37] analyze the limited capacity of dual encoders, and propose a combination of sparse and dense retrieval methods with multi-vector encoding 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. [50] proposed the DSI (differentiable search index) model to generate relevant document identifiers directly corresponding to the query. Bevilacqua et al. [4] employed the autoregressive model generate relevant words for query and utilize the generated string to retrieve relevant documents. Besides, the Deep Retrieval (DR) [20] approach for recommendation is also related this category, which learns deep retrievable network with user-item clicks and gets rid of the ANN algorithms hased on the Fuclidean space assumption.
Pre-trained language models. Recently, pre-trained Language Models (LMs), such as BERT [14] and RoBERTa [35], 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. [16] found that 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 [54] leverages hard negatives to improve the effectiveness of contrastive learning, which generates better text representations for the retrieval tasks.
3 Neural corpus indexer
The neural corpus indexer (NCT) is a sequence-to-sequene neural network model. The model takes a query as input and outputs the most relevant document identifier (docid), which can be trained by large collection of <query, docid> pairs. The documents are encoded into semantic
3

---

quary: Foods supplements liced
Generated
Encoder
PAWA Decoder
sugar?
Besearcher creammn
Ground Truth Docid
blood leve naturally when taken diily..
Query
Goouné-Totth Geneuted Qvery
Generation
Cross- Eetr-trrroy Loss
Consistency based Regularization Loss
Docidh 705
Dodid: 205
(o) Proprocessing
Query Generation
(4) Training pipaline of Neural Corpus Indexer
Figure I: Overview of Neural Corpus Indexer (NCI). (a) Preprocessing. Each document is represented by semantic identifier via hierarchical k-means. (h) Query Generation. Queries are generated for each document based the content. (c) The training pipeline of NCI. The model is trained over augmented <query, docid> pairs through a standard transformer encoder and the proposed Prefix-Aware Weight-Adaptive (PAWA) Decnder.
docids by the hierarchical k-means algorithm [23], 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 -o-uuueece transformer model [52] that takes as input the document terms and produces query as output [41]. The encoder, following the standard transformer architecture, is composed of M1 stacked transformer blocks, which outputs the representation for an imput query. For the decoder network, we stack M2 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.
3.1 Representing document with semantic identifiers
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 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 (a), given a collection of documents to be indexed, all documents are first classified into k clusters by using their representations encoded by BERT [14]. For cluster with more than e documents, the k-means algorithm is applied recursively. For each cluster containing c documents or less, each document is assigned number starting from 0 to at most c-1. In this way, we organize all documents into tree structure 7 with root rg. Each document is associated with one leaf node with deterministic routing path {ro, Ti,... Tm from the root, where r; E [0, k) represents the internal cluster index for level i, and I'm E [0, c) is the leaf node. The semantic identifier for 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 30 and e = 30 in all experiments, leaving the optimization to future work. The detailed procedure of hierarchical k-means will be described in Algorithm in the Appendix B.2.
of these
3.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, generate bunch of queries with query generation module and bind the information of document content through training the squeecce-t-sequenee model with generated queries and their corresponding document identifiers. In NCI, utilize two kinds of augmented queries:
DocT5Query. We adopt a standard sequence-to-seqquence transformer [52] based on the implementa- tion of DocT5Query [1] pre-trained by large query- -ocyumd 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 [50], 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.3 Prefix-aware weight-adaptive decoder
Given an input query æ, the probability of generating document identifier can be written as:
m
p(1 2. 0)
II prrll, F1, r2. 0z):
where is the i-th token in the current identifier; x is the representation output from encoder; 0 denotes the total parameters and 0, is the parameter for the i-th step.
This probability can be modeled by a transformer-basee decoder. For internal node with level the probabil- ity is calculated by:
12 (hf)t
h4 TransformerrDecodeerf.z h1, h2, ba......h.. h4-1; 0z),
Weight-Adopivve Projection
(2)
p(rx|ar, Pi, Softmax(h,,W) W). (3)
Here h; is the hidden representation for step i, which is calculated by multi-head attention over encoder representation x and token representations of previous decoding steps. The linear classification weight is de- noted by W E Rax=, d is the hidden dimension size and v is the vocabulary size of identifiers.
Decoder Blacks
Awwre Blocks
1
8
hg
8 F2
Position aware Embuddings
Figure 2: Overview of the Prefix-Aware Weight- Adaptive (PAWA) Decoder.
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 "52" and "5g" of the same identifier "3 "31525"" 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 "1, 1253" and "2, 4253", ",, the same token "5z" has different semantics in two different identifiers, as they are routed from different prefix paths. These two properties of the hierarchical semantic identifiers motivate to design the novel Prefix-Aware Weight-Adapte (PAWA) decoder.
Unlike standard transformer decoder, the probabilities at different tree levels, such as 0.) and p(ry (1,71.1-11.0 O,) where # 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 (2,,3) (3, 5)" for the semantic identifier "31525s", while "(2, 5(" and "(3,5) 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 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 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,
Wid AdaptiveDecoder((: F1, r2...., W;;
(4)
where e is the query embedding vector taken as initial input to the transformer decoder; frelt e (1, (1,2,..,;-111 })} are prefix tokens before the i-th position, Adeptive stacks My transformer de- coding layers with dimension d, and W E Rdxx is the adapted weight matrix for the corresponding classifier. Finally, the i-th token in the given prefix can be predicted by Softmax (h, ).
For instance, to predict the third tokens in the identifiers "(1,3) and "(),5)(,,4)3355), 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-awe aware adaptive weights can be calculated and trained in parallel in different positions while adding little burden to the entire model.
5

---

3.4 Training and inference
Consistency-basee regularization. To alleviate over-fitting, we employ a consistency-based regularization loss for training each decoding step. Given input query q. we denote the decoder representations by two forward passes with independent dropouts before Softmax as 2,,1 (r, Ti...) 11,,....--1, 02) and 2 (rr: (1. ri.....-1,,,, 0z), respectively, where B-)) de- notes the encoder network and D.)) denotes the decoder network. The consistency-baa -aaeedd regulariza- tion loss tries distinguish the representations from the same token from those of other tokens, like contrastive learning [8]. The regularization loss of query q for the i-th decoding step is defined as,
exp (sim 71,2)TT)
Lreg
log
(S)
k,p exp ( 2,x)
where we leverage dot- lot-product for sim(-); (:): Q the number of queries in the batch, and the temperature parameter is set as r = 1 in all the experiments.
Training loss. Given a of training examples D {(q, J composed of queries (training queries and augmented queries) and document identifiers, the loss function can be written as follows:
[())=
Σ (g.0)ED
(lo E(g), 0) + aCrea):
(6)
where p(d]E(() denotes the probability of generating d with & 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 sequence-to- sequence neural network, which can be optimized end-to-end via gradient descent. The hyper- parameter a denotes 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 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 (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 balanced tree structure built by corpus with N documents, the average time complexity for beam search is O N)o 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 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.
4.1 Datasets & evaluation metrics
Datasets. We conduct our experiments on two popular benchmarks for document retrieval, i.e., the Natural Questions [32] and TriviaQA dataset [29]. Natural Questions (NQ) [32] was introduced by Google in 2019. The version we use is often referred to as NQ320k, which consists of 320k: 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 reading comprchension dataset [29], which includes 78k qucry-do pairs from the Wikipedia domain. Unlike the NQ320k dataset, 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- R-Precision 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.

---

4.2 Implementation details
Hierarchical semantic identifier. For semantic identifiers, we apply a hierarchical h-means algorithm over the document embeddings obtained through a a 12-layers BERT model with pre-trained parameters (provided by HuggingFace [53]). For each hierarchical layer, we employ the default k-means algorithm implemented in scikit-learn [42] with & = 30. For simplicity, the recursion terminal condition is also set as c = 30.
Query generation. We leverage the pre-trained model, DocT5yuery [40], for query generation. We provide all document contents in NQ320k 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, 10, PyTorch 1.8.1 and HuggingFace transformers 3.4.0. We utilize the parameters of the T5 pre-trained model [5] to initialize the encoder and randomly initialize the PAWA decoder. All NCI experiments are based on learning rate 2 x 10-1 for the encoder and x 10-4 for the decoder with batch size 16 per GPU. We set the scaling factor of the consistency- baaed regularization loss as a 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 V V100 GPUs with 32GB memory. Each job takes 8 GPUs, resulting in a total batch size of 128 (16 X 8).
Baselines. We evaluate BM25 on both raw documents and those augmented by DocTsuueryy by an open-source implementation [2]. The performance of DSI [49] 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 [4] and ANCE [54] by their official implementations. Some baselines for the TriviaQA dataset are directly referred from [58]. We leave the detailed settings in Appendix B.4.
4.3 Results
In Table 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 NQ320k and TriviaQA datasets. Furthermore, an ensemble of five NCI models also brings large enhancement, because each model is trained individually with 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 perfor- mance. 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 NQ320k dataset. We find that the generated queries have different distributions with the training queries, so we also fine-tune Doc2Query on this dataset for 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% 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 based regularization. We also notice that query generation plays key role in boosting the retrieval performance. With query generation, the BM25 DocT5Query method the NQ320k dataset. ANCE
achieves higher performance than the vanilla B.M25, especially 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 retrivve--ve pipeline in traditional web search engines.
Furthermore, to study the effect of each component, we report ablation results on both NQ320k: 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 Doc Douryy This configuration removes the training queries generated by DocTSQuery. According to the results, the query generation model greatly boosts the performance. The result is
7

---

Table I: Performance comparison NQ320k retrieval task. The settings with qg-ft refer to query generation by the DocTTSuuey model fine-tuned this dataset. Other settings the original checkpoint of DoeTSSQuer..
Method
Recall
Recall
Recall
MRR @ 100
Neural Corpus Indexer (Base)
65.86
85.20
92.42
73.12
Neural Corpus Indexer (Large)
66.23
85.27
92.49
73.37
Neural Corpus Indexer (Ensemble)
70.46
94.75
77.82
Neural Corpus Indexer w/ 98 (Base)
68.91
94,48
76.17
Neural Corpus Indexer w/ 9g-fi (Large)
68.65
88.45
94.53
76.10
Neural Corpus Indexer qg-ft (Ensemble)
72.78
91.76
96.22
80.12
DSI (Basc) [50]
27.40
56.60
DSI (Large) [50]
35.60
62.60
DSI (XXL) [50)
40.40
70.30
SEAL (Base) [4]
56.98
79.97
91.39
65.48
(agge [4]
59,93
81.24
90.93
67.70
ANCE (First) [54]
51.33
80.33
91.78
61.71
ANCE (MaxP) [54]
52.63
80.38
91.31
62.84
BERT + BruteForce [15]
28.65
53.42
73.16
36.60
BERT + ANN (Fais)) [28]
27.92
53.63
73.01
37.08
BM25 DocTSQuery [40]
35,43
76.92
BM25 [44]
15.11
32.48
50.54
21.07
Table Performance comparison on TriviaQA retrieval task. The results annotated by are taken from [58].
Method
Recall@5
Recall 20
Recall@ 100
R-Precision
Neural Corpus Indexer (Base)
90.49
94.45
96.94
73.90
Neural Corpus Indexer (Large)
91.73
95.17
97.44
74.94
Neural Corpus Indexer (Ensemble)
94.60
96.89
98.20
80.84
SEAL (Base) [4]
86.3
90.5
91.5
68.1
SEAL (Large) [4]
87.7
91.8
92.6
69.2
AR2-G" (58)
78.2
84.4
87.9
coCondenser [18]
76.8
83.2
87.3
Condenser [17]
81.9
86.2
Individual Top-k [466]
76.8
83.1
87.0
Joint Top-k" [46]
RDR"(SS]
82.5
87.3
ANCE' [54]
80.3
85.3
DPR'T [30]
79.3
84.9
GAR [399]
73.1
80.4
85.7
BM25 DocT5Query [40]
59.71
72.06
82.71
39.66
BM25 [44]
56.91
69.45
80.24
37.29
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 DSI [49], using the document contents 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-adaptiv decoder has noticeable influence the performance, which indicates that, instead of borrowing the vanilla transformer decoder, it necessary to design tailored decoder architecture for the task of semantic identifier generation.
Table 3: Ablation Study on NQ320k and TriviaQA retrieval task.
Method
NQ3200
TriviuQA Recall@220 Recall@ 100) R-Precision
Recall@
Recall 10
Recall@ 100
MRR@ 100
Recall@s
Neural Corpus Indexer (Base)
65.86
85,20
92.42
73.12
90.49
94,45
96.94
73,90
w/o DocTSQuery
60.23
80.20
90.92
67.89
84.56
90.94
95.32
63.50
document query
62.49
81.21
88.85
69.41
85.34
91.10
94.66
67.48
who PAWA decoder
63.36
83.06
91.47
70.56
88.75
93.56
96.18
71.81
semantic
70.43
93.07
72.57
wfo regularization
65.07
82.91
90.65
71.80
89.01
93.63
96.16
71.59
w/o coustrained beam search
65.65
84,89
92.23
72.79
89.58
93.97
96.61
72.51

---

0.9
0.6
Receli81
0.6
0.3
0.5
02
Smell
0.4
Small
B35c
5050
2000
4000
6000
8000
1000
2000
3000
5000
Step
Step
Figure 3: Learning curves of NCI with different model capacities. Left: NQ320k: Right: TriviaQA.
w/o semantic id. This configuration replaces the semantic identifier of each document to 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 performance drop on all four metrics without using consistency-base regularization loss. The reason is that the decoder network is prone 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 tree-based prior structure. Instead, all tokens in the vocabulary can be generated in each decoding step. We observe performance drop on four evaluation metries. This indicates that it is difficult remember all information of valid identifiers in the network, and an explicit prior could be helpful for improving the quality of beam search.
Table 4: NCI with different number of layers in PAWA adapter. Left: NQ320k; Right: TriviaQA.
Setting
Recall@1
@1
Recall@ 100
6RR 100
Setting
Recall
Recall 20
Recall@ 100 R-Precision
Plaer 0 63.36
$3.06
91.47
70,56
Hlayer =
88.75
93.56
96.18
71.81
Player =
64.85
84.71
91.49
71,42
#layer
93,90
96.58
#layer 2 65.40
92.82
#layer =
89.89
94.35
96.86
72.89
Mlayer 4
65.86
85.20
92.42
73.12
Wlayer 4
90.49
94.45
96.94
73.90
Mlayer 6 65.07
83.91
91.65
71,80
#layer =6
89.76
94.31
96.76
73.32
#layer
63.60
71.22
#layer 87.90
93.30
96.20
4.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 [43]. 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 critical impact on the retrieval performance, and the capacity base model seems to be enough memorize all documents in NQ320k: 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 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 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, adopt the PAWA decoder with a 4-layers adapter in NCI.
four metrics. But when the number of layers achieves 6, the performance
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-/, A-2, B-/ and B-2, where two queries inside the same group are semantically similar, and the queries in different groups correspond distinct topics. In Figure
show the probabilities of retrieved documents

---

Lod
2-L Mad LAMAE
0412
640
0408
500 100 40% AO ANL GAN GAC
401 744
Na0
TMP 24e 201
504 1942
Figure 4: Analyses of retrieved documents with semantic identifiers. Left: The probabilities of retrieved documents for Query Group A: Middle: The probabilities for Query Group B; Right: The 1-SNE visualization of BERT-based document embeddings.
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 1-SNE [51] in Figure 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, conclude that the semantic document identifiers generated by the hierarchical k-means algorithm have positive effects on the retrieval performance.
Efficiency Analysis. We use 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 end-to-end retrieval method and achieves competitive performance without re-ranking. the latency and throughput 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 [2]), 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.
Table 5: Efficiency analysis
Model
Beam
Latency
Throughput
size
size
(ms)
(queries/ /s)
Small
10
78.46
58.48
Base
115.17
52.55
Large
10
188.60
43,39
already affordable for some
Small
100
216.01
6.12
269.31
Large
100
356.07
4.75
5 Limitation & Future Works
Despite the significant breakthrough, the current implementation of NCI still suffers from several limitations before deployment in 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 update the model-basee 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) [47] 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. this way, all documents in relevant clusters can be retrieved efficiently. (3) Model compression techniques, like weight quantization [26] and knowledge distillation [24], 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 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- -oceumet pairs. To optimize the recall and ranking performance, we invent a tailored prefix-aware aeight-adaptive decoder. Empirically, we evaluate NCI on NQ320k and TriviaQA datasets, demonstrating its outstanding performance over tfte-off-hhe-rt solutions.
10

---

References
[2]
[3] Jon Louis Bentley. Multidimensional binary search trees used for associative searching. Communications of the ACM, 18(9):50--517, 1975.
[4] Michele Bevilacqua, Giuseppe Ottaviano, Patrick Lewis, Wen-tau Yih, Sebastian Riedel, and Fabio Petroni. Autoregressive search engines: Generating substrings as document identifiers. arXiv preprint arXiv:2204.1065 10628, 2022.
[5] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
[6] Wei-Cheng Chang, X Yu Felix, Yin-Wen Chang, Yiming Yang, and Sanjiv Kumar. Pre-training tasks for embedding- based large-scale retrieval. In International Conference on Learning Representations, 2019.
[7] Qi Chen, Bing Zhao, Haidong Wang, Mingqin Li, Chuanjie Liu, Zengzhong Li, Mao Yang. and Jingdong Wang. Spann: Highly-efficient billion-scale approximate nearest neighbor search. arXiv preprint 1...:2111.08866. 2021.
[8] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pages 1597-1607. PMLR, 2020.
[9] Zhuyun Dai and Jamie Callan. Context-aware sentence/passage term importance estimation for first stage retrieval. arXiv preprint 191iv:1910.10687 2019.
[10] Mayur Datar, Nicole Immorlica, Piotr Indyk, and Vahab S Mirrokni. Locality sensitive hashing scheme based p-stable distributions. In Proceedings of the twentieth annual symposium Computational geometry, pages 253-262, 2004.
[11] Nicola De Cao, Wilker Aziz, and Ivan Titov. Highly parallel autoregressive entity linking with discriminative correction. arXiv preprint NSKKi:2119.03777 2021.
[12] Nicola De Cao, Gautier Izacard, Sebastian Riedel, and Fabio Petroni. Autoregressive entity retrieval. arXiv preprint arXiv:2010.00904, 2020.
[13] Nicola De Cao, Ledell Wu, Kashyap Popat, Mikel Artetxe, Naman Goyal, Mikhail Plekhanov, Luke Zettlemoyer, Nicola Cancedda, Sebastian Riedel, and Fabio Petroni. Multilingual autore- gressive entity linking. arXiv preprint arXiv:2103. 12528, 2021.
[14] Jacob Devlin, Ming- Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint trXiv:1810.044 2018.
[15] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 4171 4186, 2019.
[16] Luyu Gao and Jamie Callan. Condenser: a pre-training architecture for dense retrieval. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 981-993, 2021.
[17] Luyu Gao and Jamie Callan. Is your language model ready for dense representation fine-tuning. arXiv preprint SoXi::21040082:: 2021.
[18] Luyu Gao and Jamie Callan. Unsupervised corpus aware language model pre-training for dense passage retrieval. arXiv preprint 10t.:22088.0554400 2021.

---

[19] Luyu Gao, Zhuyun Dai, and Jamie Callan. Coil: Revisit exact lexical match in information retrieval with contextualized inverted list. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, pages 3030-3042, 2021.
[20] Weihao Gao, Xiangjun Fan, Chong Wang, Jiankai Sun, Kai Jia, Wenzhi Xiao, Ruofan Ding, Xingyan Bin, Hui Yang, and Xiaobing Liu. Deep retrieval: Learning a retrievable structure for large-scale recommendations. arXiv preprint arXiv:2007.07200 2020.
[21] Jiafeng Guo, Yixing Fan, Qingyao Ai, and W Bruce Croft. A deep relevance matching model for ad-hoc retrieval. In Proceedings the 25th ACM international on conference on information and knowledge management, pages 55 64, 2016.
[22] Tonglei Guo, Jiafeng Guo, Yixing Fan, Yanyan Lan, Jun Xu, and Xueqi Cheng. comparison between term-based and embedding-basse methods for initial retrieval. In China Conference on Information Retrieval, pages 28-40. Springer, 2018.
[23] John A Hartigan and Manchek A Wong. Algorithm as 136: A k-means clustering algorithm. Journal of the royal statistical society. series (applied statisties), 28(1):100-108, 1979.
[24] Geoffrey Hinton, Oriol Vinyals, Jeff Dean, al. Distilling the knowledge in neural network. arXiv preprint arXiv: PeXiv:1500 150111003..2222 02552, 2(7), 2015.
[25] Po-Sen Huang, Xiaodong He, Jianfeng Gao, Li Deng, Alex Acero, and Larry Heck. Learning deep structured semantic models for web search using clickthrough data. In Proceedings of the 22nd ACM international conference on Information & Knowledge Management, pages 2333-2338, 2013.
[26] Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient aiitgge--rrthmeetic--e
inference. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2704 2713, 2018.
[27] Suhas Jayaram Subramanya, Fnu Devvrit, Harsha Vardhan Simhadri, Ravishankar Krishnawamy, and Rohan Kadekodi. Diskann: Fast accurate billion-point nearest neighbor search on single node. Advances in Neural Information Processing Systems, 32, 2019.
[28] Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, 7(3):535 547, 2019.
[29] Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint 1xXv::1705.03551, 2017.
[30] Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. arXiv preprint arXiv: 2ixx:2004.04906 2020.
[31] Omar Khattab and Matei Zaharia. Colbert: Efficient and effective passage search via contextual- ized late interaction over bert. In Proceedings of the 43rd International ACM SIGIR conference on research and development in Information Retrieval, pages 39 48, 2020.
[32] Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453-466, 2019.
[33] John Lafferty and Chengxiang Zhai. Document language models, query models, and risk minimization for information retrieval. In Proceedings of the 24th annual international ACM SIGIR conference on Research and development information retrieval, pages 111-119, 2001.
[34] Canjia Li, Andrew Yates, Sean MacAvaney, Ben He, and Yingfei Sun. Parade: Passage representation aggregation for document reranking. arXiv preprint arXiv:2008.0909 2020.
12

---

[35] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer. and Veselin Stoyanov. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/1907.1 1692, 2019.
[36] Wenhao Lu, Jian Jiao, and Ruofei Zhang. Twinbert: Distilling knowledge twin-structured compressed bert models for large-scale retrieval. In Proceedings of the 29th ACM International Conference Information & Knowledge Management, pages 2645-2652. 2020.
[37] Yi Luan, Jacob Eisenstein, Kristina Toutanova, and Michael Collins. Sparse, dense, and attentional representations for text retrieval. Transactions of the Association for Computational Linguistics, 9:329-345, 2021.
[38] Yu A Malkov and Dmitry A Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE transactions on pattern analysis and machine intelligence, 42(4):824-83) 836, 2018.
[39] Yuning Mao, Pengcheng He, Xiaodong Liu, Yelong Shen, Jianfeng Gao, Jiawei Han, and Weizhu Chen. Gennratioon-augmenndd retrieval for open-domain question answering. arXiv preprint arXiv:2009.08553, 2020.
[40] Rodrigo Nogueira, Jimmy Lin, and Al Epistemic. From doc2query to doctttttquery. Online preprint, 2019.
[41] Rodrigo Nogueira, Wei Yang. Jimmy Lin, and Kyunghyun Cho. Document expansion by query prediction. arXiv preprint arXiv: 1eXiv:190400837. 2019.
[42] Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, et al. Scikit- learn: Machine learning in python. the Journal of machine Learning research, 12:2825-2833 2011.
[43] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 1910. 10683, 2019.
[44] Stephen Robertson and Hugo Zaragoza. The probabilistic relevance framework: BM25 and beyond. Now Publishers Inc, 2009.
[45] Stephen E Robertson and Steve Walker. On relevance weights with little relevance informa- tion. In Proceedings of the 20th annual international ACM SIGIR conference on Research and development in information retrieval, pages 16 16-24, 1997.
[46] Devendra Singh Sachan, Mostofa Patwary, Mohammad Shoeybi, Neel Kant, Wei Ping, William L. Hamilton, and Bryan Catanzaro. End-to-end training of neural retrievers for open- domain question answering. arXiv preprint arXiv:2101.00400: 2021.
[47] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. International Conference on Learning Representations (ICLR), 2017.
[48] Yelong Shen, Xiaodong He, Jianfeng Gao, Li Deng, and Grégoire Mesnil. Learning semantic representations using convolutional neural networks for web search. In Proceedings of the 23rd international conference on world wide web, pages 373-374, 2014.
[49] Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention in transformer models. arXiv preprint arXiv:2005.00744 2020.
[50] Yi Tay, Vinh Q Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Gupta, et al. Transformer memory as a differentiable search index. arXiv preprint arXiv:2202.0699
2022.
[51] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal machine learning research, 9(11), 2008
13

---

[52] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS), pages 5998-6008, 2017.
[53] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et Huggingface's s transform- ers: State-of-the-arr natural language processing. arXiv preprint 1910:0977.037,,, 2019.
[54] Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul N Bennett, Junaid Ahmed, and Arnold Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval. In International Conference Learning Representations, 2020.
[55] Sohee Yang and Minjoon Sco. Is retriever merely approximator of reader? arXiv preprint arXiv:2010.1 10999, 2020.
[56] Wei Yang. Haotian Zhang, and Jimmy Lin. Simple applications of bert for ad hoc document retrieval. arXiv preprint arXiv: 1901. 19972, 2019.
9033
[57] Hamed Zamani, Mostafa Dehghani, W Bruce Croft, Erik Leamed-Miller, and Jaap Kamps. From neural re- re-ranking to neural ranking: Learning sparse representation for inverted index- ing. In Proceedings of the 27th ACM international conference information and knowledge management, pages 497-506, 2018.
[58] Hang Zhang, Yeyun Gong, Yelong Shen, Jiancheng Lv, Nan Duan, and Weizhu Chen. Ad- versarial retriever-ranker for dense text retrieval. In International Conference on Learning Representations, 2021.
[59] Shengyao Zhuang, Hang Li, and G. Zuccon. Deep query likelihood model for information retrieval. In ECIR, 2021.
Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information on how to answer these questions. For each question, change the default [TODO] to [Yes], [No], or [N/A] You are strongly encouraged to include a justification your answer, either by referencing the appropriate section of your paper or providing brief inline description. For example:
Did you include the license to the code and datasets? [Yes] See Section.
Did you include the license the code and datasets? [No] The code and the data proprictary.
Did you include the license to the code and datasets? [N/A]
Please do not modify the questions and only the provided macros for your answers. Note that the Checklist section does not count towards the page limit. In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section 5.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See appendix.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you experiments...
14

---

(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Section 4.
(c) Did you report error bars (e.g., with respect the random seed after running experi-
ments multiple times)? [No]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Section
4. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]] See Section 4.
(b) Did you mention the license of the assets? [Yes] See appendix.
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [Yes]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [Yes]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
15


Reviews:
[{"review_id": "xSA8ciX29e1", "reviewer": "Reviewer_xW4j", "summary": "\nThis paper proposes a new framework for neural IR: given query, directly predict a document ID. The document IDs are obtained by hierarchical clustering of documents beforehand. This is a novel formulation of the problem, and is very distinct from current two-stage methods that have a high-recall sparse retrieval stage, followed by a high-precision neural reranker, or approximate nearest neighbor methods that encode both documents and queries as vectors.  ", "questions": "- Can you provide some timing analysis for the baselines (e.g. in footnote), so Table 4 has some reference of comparison?", "limitations": "- The limitation section is reasonable. I think the index-update aspect may be one of the most important near-term challenges. ", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\nStrength:\n- The formulation is unique and intriguing. I really like this idea! Despite some limitations, the results are very promising and I imagine many follow-up work may come from this paper. \n\nWeakness:\n- I think the method rests on one large assumption, that documents can be pre-clustered into some fixed set of IDs. All queries are being trained with the same fixed set of document IDs. This is a somewhat unconventional way to think about IR, as one usually does not assume document cluster structure, and ad hoc queries may flexibly retrieve any subset. I think some experiments that try different definitions of document IDs (e.g. random) would shed light on the assumption of having a strong document clustering prior. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "ZFdyQa4Tae2", "reviewer": "Reviewer_V4yZ", "summary": "While the traditional document retrieval problem is formulated as mapping both the document and the query to the same vector space and then performing nearest neighbor search to find the closest document fore each query, the paper instead proposes to formulate the task as a document identifier generation problem, where the input is the query and the output is the target document's id. This is not entirely novel with respect to DSI (Tay et al., 2022) but given that DSI was available on the archive only three months before NeurIPS deadline, one might be able to consider this paper to be a concurrent work. The proposed model by this paper, Neural Corpus Indexer, also introduces a few techniques: (1) query generation that augments the query-to-doc data for training (equivalent to document indexing), (2) prefix-aware weight-adaptive decoder that gives different symbols to the same tokens in different positions of the ids, (3) semantic document id so that similar documents share similar ids, and (4) regularization technique to prevent overfitting. The results on Natural Questions 320k (a subset of NQ) are very promising; the proposed model not only outperforms other generative retrieval models such as DSI and SEAL by a big margin but also strongly outperforms competitive baselines such as ANCE and BM25. ", "questions": "- With regard to the weakness mentioned above, I am wondering how sensitive the model's performance is with respect to the quality of the generated queries.\n- I don't think it is correct to say that \"DSI hardly captures the hierarchical semantics of document identifiers\" in L42. DSI indeed proposes a similar clustering mechanism to what this paper proposes.\n- How is the documents clustered for semantic identifiers? Could you please give more description of Section 3.1.\n- What is difference between p1 and p in L223? I think Section 3.4 is unclear.", "limitations": "Yes, the paper mentions crucial limitations such as its lack of scalability compared to canonical NNS retrieval system, which I agree as well.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The results are very strong. The baselines that the paper compares against, such as ANCE and BM25, are very strong ones and hard to beat. It is remarkable that a generative model can beat those baselines.\n- Query generation helps a lot and this is a very useful information for people working in this field. In fact, as shown in Table 2, among other techniques introduced by the paper, QG is by far the biggest factor for the superior performance of NCI.\n\nWeaknesses:\n- If query generation is the dominant factor, then it seems to me that what the model is really doing is \"query indexing\" rather than \"document indexing\". This means the performance of the model will highly depend on generating a comprehensive range of queries which are likely to overlap with the test queries. \n- Constructing the semantic hierarchy of the document ids can be considered as constructing the inverted index in the canonical NNS retrieval setup. Then what the decoder is doing can be considered as locating the correct cluster in each level of the hierarchy. Then I am wondering if it is fair to say this model is really \"generative retrieval\" or it is more of using different embeddings for the nearest neighbor search in each layer of the hierarchy. \n- The paper only experiments on Natural Questions and only a subset of it. While NQ is a very representative dataset in this task, the paper would have looked much better if it was tested on different datasets and/or the full dataset of NQ. It is not clear why such results are not shown.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Yz8jrsogSL", "reviewer": "Reviewer_ncpT", "summary": "The paper proposed a sequence-to-sequence model that generates the relevant document given an input query, as a retrieval method. The proposed method have a huge improvement on the NQ320k dataset (Natural Questions). The most improvement seems to be coming from using generated query (based on the document) paired with the document as augmented training data.", "questions": "1. Can we use the same k-means procedure to obtain an identifier for a new document? I guess new documents will have the same identifier to be meaningful to the model?\n2. It seems that the query generations has most contribution to the improvement. Have you compared with other query generation models / techniques so that we can understand what contribute to the success of this method?", "limitations": "The authors acknowledge a set of limitations of the model, including scale of the index set, inference speed and updating the index. I'm most interested in the third one. The model can only work on a fixed set of documents (to retrieve from), or there is a cost to update the model when the index set changes. This is a limitation of the line of works that generating the document identifiers, not specific to this work. However, I'm also curious, given that the document identifier is generated using a k-means algorithm, how the model performs when there are new documents if we just assign an identifier to a new document using the same k-means algorithm.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality:\nThe proposed data augmentation technique is interesting and demonstrate good performance. The other techniques are interesting and providing much smaller but also good gains.\n\nQuality:\nThe author addressed their limitations and present the work in context. The experiments are well designed.\n\nClarity:\nThe paper is clearly written.\n\nSignificance:\nVery good gains on NQ320k over previous best method.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "5T633xe-7i3", "reviewer": "Reviewer_RUss", "summary": "The paper focuses on a promising direction to learn retrieval via generation, where the model will generate the relevant document IDs for given queries. It resembles the previous Differentiable Search Index work with several minor improvements.\n", "questions": "NA", "limitations": "NA", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nThey propose several techniques to improve the model such as prefix-aware weight-adaptive decoder, query generation, consistency-based regularization. They show excellent improvement in recall@1 and recall@10 on the NQ320k datasets, largely outperforming previous baselines.\n\nWeaknesses:\n\nNCI is quite similar to DSI - e.g., the encoder-decoder architecture, indexing training and the hierarchical clustering algorithm to generate semantic doc IDs. It introduces several small improvements such as using prefixLM, consistency-based regularization, query generation, etc. The overall contribution seems limited. \n\nAnother weakness is that many of the gains come from using the additional query generation data. Interestingly, all other improvements seem to have very limited effect on the performance. This also raises a question about how much the model relies on the quality of the query generation model and how it can adapt to new domains that the QGEN model can’t generate good questions for the target document. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept