Transformer Memory as a Differentiable Search Index
Authors: Yi Tay
Vinh Q. Tran
Mostafa Dehghani
Jianmo Ni
Dara Bahri
Harsh Mehta
Zhen Qin
Kai Hui
Zhe Zhao
Jai Gupta
Tal Schuster
William W. Cohen
Donald Metzler
Venue: NeurIPS 2022 Accept (2022)

Abstract: In this paper, we demonstrate that information retrieval can be accomplished with a single Transformer, in which all information about the corpus is encoded in the parameters of the model. To this end, we introduce the Differentiable Search Index (DSI), a new paradigm that learns a text-to-text model that maps string queries directly to relevant docids; in other words, a DSI model answers queries directly using only its parameters, dramatically simplifying the whole retrieval process. We study variations in how documents and their identifiers are represented, variations in training procedures, and the interplay between models and corpus sizes. Experiments demonstrate that given appropriate design choices, DSI significantly outperforms strong baselines such as dual encoder models. Moreover, DSI demonstrates strong generalization capabilities, outperforming a BM25 baseline in a zero-shot setup.

Transformer Memory as a Differentiable Search Index
Yi Tay", Vinh Q. Tran", Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta
Zhen Qin, Kai Hui, Zhe Zhao, Jai Gupta, Tal Schuster
William W. Cohen, Donald Metzler Google Research {yitay, mgtrann,.etzleerraggggl
Abstract
In this paper, demonstrate that information retrieval can be accomplished with a single Transformer, in which all information about the corpus is encoded in the parameters of the model. To this end, we introduce the Differentiable Search Index (DSI), a new paradigm that learns text-to-text model that maps string queries directly to relevant docids; in other words, DSI model answers queries directly using only its parameters, dramatically simplifying the whole retrieval process. We study variations in how documents and their identifiers are represented, variations in training procedures, and the interplay between models and corpus sizes. Experiments demonstrate that given appropriate design choices, DSI signifi- cantly outperforms strong baselines such as dual encoder models. Moreover, DSI demonstrates strong generalization capabilities, outperforming a BM25 baseline zero-shot setup.
Introduction
Information retrieval (IR) systems map user query q Q ranked list of relevant documents (di...., dn C DD typically represented by integers or short strings called document identifiers based on pipelined retrieve-then-rank strategies.
(docids). The most widely used IR approaches
For retrieval, approaches based on inverted indexes or nearest neighbor search are common where contrastive learning based dual encoders (DEs) (Gillick et al. 2018 Karpukhin et al. 2020 Ni et al. 2021) are the present statc-of-the-art.
This paper proposes alternative architecture, wherein a system (Sutskever et al., 2014) is used to directly map a query q to a relevant docid j E Y. This proposal is shown in the bottom half of Figure [] for architecture.
(seq2seq) learning
encoder-decoder
We call this proposed architecture a differentiable search index (DSI), and implement it with large pre-trained Transformer (Vaswani et al.) (007)) model, building on the recent success of large generative language models (LMs) (Brown al.) 2020; Raffel et al. 2019 Devlin et al.) 2018 Thoppilan et al. 2022 Du et al. 2021) In this proposed architecture, all information of the corpus is encoded within the parameters of the Transformer language model.
At inference time, the trained model takes input a text query q and outputs docid j. If desired, beam search can be used to produce ranked list of potentially-relevant docids. As we show, this process can work surprisingly well when trained properly. In our experiments can consistently outperform DE baselines, sometimes drastically: for a base-sized T5 model, Hits on the smallest corpus is improved by more than 20 points, from 12.4% for a DE to 33.9% for DSI; and on corpus
"Co-leads.
36th Conference on Neural Information Processing Systems (NcurIPS 2022), New Orleans, LA, USA.

---

Learning encode
Retrieve
query.22
quaRy1a3
Ericoder
Vizs
Who teppin
doc456
Maximaal Inner
As part of her research. Sarah Waners came soroes the thre her first book, ltpp ng the Velvet...
Encoder
Dommeni carritt
Product Search
Rank List
doc131
Eass
doo456 doc253 :
Kocley Howzs Known for her tolcas Kitty Butler in Tipping the Velvet
2
cuary:23
Leaming encode and retrieve
Bearn Search
Pank List
Who wuthor Hoppi, Ehe vilvet?
Fize
dcc456
das456 00c283
doc456
As part of her research Samhh Whiters come scroes book, Trap no Vidvot....
Encoder
Vips
Decader
doc 137
dcc131
Howes for Kitty
Tipping the
Figure Comparison of dual encoders (top) to differentiable search index (bottom).
Table I: Information retrieval requires series of decisions, associated with the subproblems of document representation, indexing, and retrieval. Structured-documment variants of DSI are also sensitive to fourth decision, namely how docids are represented.
BM25 TFIDF
Dual Encoder (DE)
Differentiable Search Index (DSI)
doc/query rep.
sparse Vd, vector in R VI
dense Vd, vector in Rd
Various (see Section 3.1.2
docid rep.
Various (see Section 3.2)
indexing
build inverted index mapping each term (d d
build table mapping each docvec Vd, + approximate sparse matmul to find argmax, v. Vd,
train model (see Section 3.1.1
to map d, +
retrieval (top-1)
approximate MIPS to find argmax, v,, Vd,
trained model to find argmax, Pr(jlg)
30x larger, performance is improved by nearly 7 points. These gains increase when larger models are used: for an LIB-pap 1B-parameter T5 model, Hits@@ performance improves by more than 25 points over DE on the small corpus, and more than 15 points on the large corpus. DSI also performs extremely well in zero-shot setting, e-g., improving Hits by 14 points over BM25.
In addition to these quantitative gains, the DSI architecture is much simpler than a DE (see Table []. DE system fixes search procedure (MIPS) and learns internal representations that optimize performance for that search procedure; in contrast, a DSI system contains no special-purposs fixed search procedure, instead using standard model inference map from encodings to docids.
Of particular interest to the machine learning community, as Table I shows, in DSI all aspects of retrieval are mapped into well-understood ML tasks. This may lead to new potential approaches to solving long-standing IR problems, As one example, since indexing is now a special case of model training, incrementally updating index becomes special case of model updating (Sun et al. 2020).
In this paper, DSI is applied moderate- sized corpora (from 10k to 320k documents), of which are derived from one challenging retrieval task, and we leave the important question of the scaling DSI to larger corpora to future work. The task considered retrieving supporting passages given questions from the Natural Questions (NQ) dataset, a challenging task for lexical models.
While the idea of DSI simple, there are number of ways can be realized, some of which work surprisingly well, and some of which work surprisingly poorly. Below we explore a number of variations of the DSI architecture.
Document representation. We explore several approaches to representing documents, including a "naive" approach of using the document's 's full text, well as variants of the bag-of-ww of-dd representa- tion used by traditional IR engines.
Docid representation. We look at several ways to represent docids. In addition to naively representing integers as text strings, we also consider unstructured atomic docids, where each document is assigned unique token, and some simple baselines for constructing structured semantic docids that describe

---

how to navigate to a document through a hierarchical clustering of the corpus. Structured docids either semantically structured via clustering, or naively structured as tokenized to large corpora, since the size of the vocabulary used in the decoder is made larger.
better
Indexing. A trainable IR system traditionally has two phases: indexing corpus (i.e., memorizing information about each document), and learning how to effectively retrieve from the index. In DSI, the index is stored in the model parameters, and indexing is simply another kind of model training. Figure D suggests one approach to indexing corpus: namely, to train on (1) examples (dj, j) that pair document d with its docid j. in addition to (2) examples (g. j) that pair query q with a relevant docid j. In this setup the examples of type (1) are "indexing" examples.
While it is clear that examples of type (2) alone do not provide enough information for system to generalize to novel retrievals, there are many alternatives to examples of type (1) that might plausibly "teach" a model about the associations between documents and docids. We explore number of these below, and show that some plausible-seeming techniques perform very poorly. We also explore a number of alternative multi-task optimization and curriculum learning schemes for combining these types of examples.
Effects of model and corpus size. Since recent results suggest that some properties of large LMs emerge only for very large model sizes Brown et al. (2020), we explore the performance of DSI for a range of model sizes and corpus sizes of I0k, 100k, and 320k documents.
Summary. We show that even naive representations for documents and docids, coupled with appropri- ate training procedures to fine-tune modern large LMs, can perform surprisingly well; we present two improved docid representations, unstructured docids and ammantically-structune improve the naive representation choice. We show that there is substantial variation in performance among indexing/training strategies and we show that performance of DSI significantly and consis- tently improves with model scale. To our knowledge this is the first case of generative indexing improving performance over strong baselines for a well-studied document retrieval task.
docids, which
2 Related Work
De Cao et al. (2020) describe a related sequence-to-sequenee system called autoregressive entity linking, which documents mentioning an
implicitly, c.g., by posing question to which that entity is an answer are mapped to a canonical name of that entity. In the case of Wikipedia, canonical entity names correspond to page titles, so this could be viewed as a sort of document retrieval. This approach has been adapted to other purposes, such as generating knowledge base triples in canonical form (Josifoski et al., 2021). The task we consider different from those considered in autoregressive entity linking: our goal is to retrieve a document containing the answer, rather than a document whose title is the answer. More importantly, in autoregressive entity linking the generation target is a semantically meaningful name, whereas we allow targets to be arbitrary docids. This makes our approach applicable to general retrieval tasks, but raises new questions about docid representation and indexing strategies.
In autoregressive entity linking, generation is constrained to return an output from a fixed set. would be feasible to constrain DSI generation outputs to be valid docids. Although we do not use this technique, the degree to which this might improve performance is a worthwhile question.
There is a large body of work on retrieval augmented generation, i.e., retrieving auxiliary documents enhance language models (Borgeaud et al. 2021 Guu al. 2020). These techniques are useful for many tasks including question-answering. but rely on traditional retrieval methods such as DEs. Here we use generation to replace a retrieval process, rather than using retrieval augment a generation process.
Dual encoders (Dehghani al. 2017 Gillick al. 2018 Gao et al. 202 Ni al. 2021
Karpukhin
et al. 2020) are a well-established paradigm for retrieval. The key idea is produce query and document embeddings independently and perform a similarity retrieval in vector space across all embedding pairs. Query and candidate documents are produced by a sequence encoder and training is performed using a form of contrastive loss.
The interpretation of a large Transformer model as a memory store have been investigated in prior work. (Roberts et al. 2020) demonstrated success a closed-book QA task whereby they train
3

---

T5 models to retrieve facts that are encoded within the parameters of the model during pretrain- ing. However, different from CBQA, the presented problem here in this paper is retrieve full documents based on docids instead of generating direct answers. Meanwhile, (Petroni et al. 2019) also investigated language models as knowledge bases and found that pretrained LMs may already contain relational knowledge. (Geva et al. (020) analyzes the knowledge encoded within Transformer feedforward layers. There have been also works that demonstrate the relation of Transformers to associative memory and Hopfield networks (Ramsauer et al. 22220), which reinforce the notion that Transformers should intuitively serve as good associative memory store or search index.
Differentiable Search Index
The core idea behind the proposed Differentiable Search Index (DSI) is fully parameterize tradi- tionally multi-stage retrieve-then-rank pipelines within a single neural model. To do so, DSI models must support basic modes of operation:
Indexing: DSI model should learn associate the content of each document d, with its corresponding docid j. This paper utilizes straightforward lequence-to-sequeee (seq2seq) approach that takes document tokens input and generates identifiers output.
Retrieval: Given imput query, DSI model should return ranked list of candidate docids. Here, this achieved with autoregressive generation.
Following these two operations, DSI model can be trained to index corpus of documents and optionally fine-tune
available set of labeled data (queries and labeled documents), and thereafter used to retrieve relevant documents- all within a single, unified model. As opposed to retrieve-then- rank approaches, this type of model allows for simple end-to-end training and can easily be used as differentiable sub-component larger, more complex neural model.
3.1 Indexing Strategies
We investigate various indexing strategies that are meant to learn associations between documents and their identifiers. We train our model to predict docids given a sequence document tokens. This allows model to learn which identifier belongs to which document and can be thought of as differentiable take on traditional search indexes. We consider various alternatives and ablate these settings in subsequent sections. The final strategy employed was Inputs2Targets with direct indexing.
3.1.1 Indexing Method
This section discusses the indexing task variants that we consider.
-> docid. As its suggests, this binds the docids to the document tokens in straightforward inputs-to-targets fashion. The advantage here is that the identifier is the denoising target, which puts closer proximity the loss function. Since the retrieval task is also concerned with predicting identifiers, this formulation allows the network to follow a similar input-target balance in terms of sequence length. A potential weakness is that the document tokens are not denoising targets and therefore there is no opportunity for general pre-traiee pre-training on document tokens.
Inputs2Target We frame this seq2seq task of
This formulation considers the opposite of the above. i.e., generating document tokens from identifiers, ie.. docid doc doc_ toketos.ss. Intuitively, this is equivalent training an autoregressive language model that is conditioned on the docid.
Targets2Inputs
Bidirectional
This formulation trains both Inputs2Targeets and Targets2Inputs within the same co-training setup. A prefix token is prepended to allow the model to know which direction the task is being performed in.
Span Corruption We also explored setup that performs span corruption-based denoising (Raffel al. 2019) with the inclusion of docid tokens. In this approach, we concatenate the identifier the document tokens prefix that can be randomly masked spans in the span corruption objective.

---

This method has the advantage of (I) also performing general pre- vre-training during indexing and (2) achieving good balance of docids as denoising targets and inputs.
3.1.2 Document Representation Strategies
In the previous section, we explored "how to index". This section investigates "what to index?", ie.. how to best represent toc_tokens, We state our options here and carefully ablate them in our experiments later. The best option in the end was the direct indexing method.
Direct Indexing This strategy represents a document exactly. We take the first L tokens document, with sequential order preserved, and associate them with the docid.
Set Indexing Documents may contain repeated terms and/or non-informative words (e.g., stop- words). This strategy de-duplicates repeated terms using the default Python set operation and removes stopwords from the document. The rest of the document after filtering is passed into the model similar fashion to the direct index.
Inverted Index This strategy maps chunked documents (contiguous blocks of tokens) instead of entire documents directly to the docid. We randomly subsample a single contiguous chunk of tokens and associate them with the docid. The key advantage of this approach is to allow looking beyond the first k tokens.
3.2 Representing Docids for Retrieval
Retrieval within seq2seq- based DSI models is accomplished by decoding docids given input query. How to do this decoding in an effective way largely depends on how docids are represented in the model. The remainder of this section explores number of possible ways for representing docids and how to handle decoding for each.
Unstructured Atomic Identifiers The most naive way to represent documents assign each an arbitrary (and possibly random) unique integer identifier. We refer to these as unstructured atomic identifiers.
With these identifiers, obvious decoding formulation is to learn probability distribution over the identifiers. In this case, models are trained to emit one logit for each unique docid ( ([Ndocuments)) This is analogous to the output layer in standard language models, but extended to include docids.
To accommodate this, extend the output vocabulary standard language model as follows:
0 = S-offoax [tokenss W Wdoce] T huast)
where [:] is the row-wise concatenation operator, Wtokens E [Rdmodet x Nestenal and Wdocs E Rddmoded / [Naccmmantal, Mast is the last layer's hidden state (E (Remode)) of the decoder stack. To retrieve the top-k documents for given query, we simply sort the output logits and return the corre- sponding indices. This is also reminiscent of standard listwise learning to rank where all documents are considered at once.
Naively Structured String Identifiers We also consider an ostensibly absurd approach that treats unstructured identifiers, i.e., arbitrary unique integers, as tokenizable strings. We refer to these as naively structured identifiers.
In this formulation, retrieval is accomplished by decoding a docid string sequentially one token at time. This eliminates the need for the large softmax output space that comes with unstructured atomic identifiers. It also eliminates the need to learn embeddings for each individual docid.
When decoding, beam search is used to obtain the predicted best docid. With this strategy, it is less straightforward to obtain a top-k ranking. One could exhaustively comb through the entire docid space and obtain the likelihood of each docid given the query. Instead, we use the partial beam search tree to construct top-k retrieval scores. We find this approximation be quite efficient and effective in practice.
5

---

Semantically Structured Identifiers All of the approaches for representing docids thus far as- sumed that the identifiers are assigned in an arbitrary manner. While exploring the limits of arbitrary identifiers is quite interesting, it is only intuitive that imbuing the docid space with semantic structure can lead better indexing and retrieval capabilities. As such, this section explores semantically structured identifiers.
Specifically, we aim to automatically cre- ate identifiers that satisfy the following properties: (1) the docid should capture some information about the semantics of its associated document, (2) the docid should be structured in a way that the search space is effectively reduced after each decoding step. This results in identifiers where semantically similar documents share identifier prefixes.
Cocumett
12
21
22
23
Beom search decodda 233'
Figure 2: Visual example of a hierarchical cluster- ing process used to assign semantically structured identifiers. During inference, beam search navi- gates this trie to decode the correct docid.
In this work, treat this fully unsuper- vised pre-processing step. However, as part of future work it may be possible to integrate and automatically learn semantic identifiers fully end-to-end manner.
Algorithm Generating semantically structured identifiers. (Referenced in Section 3.2
To construct identifiers with this property, we employ a simple hierarchical clustering process over document embeddings to in- duce a decimal tree (or more generally, trie).
Input: Document embeddings X N where X, E R
Output: Corresponding docid strings J.NN
function GENERATRSEMANTTON (NTICIDS(XXX
C1:10 Cluster(X J <- empty list 0 do
10)
Jeurrent < [8] *CC++1
Given corpus to be indexed, all documents are clustered into 10 clusters. Each document is assigned identifier with the number of their cluster from 0-9, For every cluster con- taining more than documents, the algorithm applied recursively, with the next level's result (the remaining suffix of the identifier) appended to the existing identifier.
if CC41 > then
Jreut -ENERATTSERANTICCDDICC
else
Jree [0..... -1
end if Jauster deceeaaaaiieeSrrrrooaaaaa Jeurvents Jreat)
- JappenddElemeesss (ettstt)) end J reando/ToOrrginnll, X1.N. C1ao)
return J function
For clusters with c documents or less, each element is assigned an arbitrary number from 0 to at most and likewise its digits appended to the existing identifier. Although this specific process induces a decimal tree, it is possible to induce similar types of tries using any number of other reasonable strategies. In practice,
simply apply k-means over embeddings generated by small 8-layer BERT model, with c = 100. We include pseudo-code for this process Algorithm
3.3 Training and Optimization
The DSI models that we train are optimized for seq2seq cross entropy loss and are trained with teacher forcing. We explored two main strategies for training DSI models. The first and more straightforward strategy first train model to perform indexing (memorization), followed by fine-tuning stage where the trained model is used to map queries to docids (e.g., retrieval). The second strategy is to train them together a multi-tase tast setup. To this end, we frame co-training tasks in similar fashion to T5-style co-training (c.g., using task prompts to differentiate them). The latter performed significantly better, especially when the proportion of indexing to retrieval task examples is high. Hence, we adopted multi-task learning as the default strategy.
Here, we make the observation that our setup unique and unlike traditional multi-task learning or transfer learning. In typical multi-task setups, two tasks have shared commonalities that could improve the performance of both tasks if they were learned together. However, in our setup, the retrieval task is completely dependení on the indexing task. In particular, without the indexing task, the identifiers leveraged by the retrieval task would be completely meaningless. Hence, in order to solve task B (retrieval), the model needs to learn task A (indexing) well enough. This problem

---

setup presents unique and largely unexplored research challenges that might be of interest to the ML community.
4 Experiments
In this section, we discuss our experimental setup, datasets used and baselines compared. We also discuss experimental results, findings and effect of various strategies discussed in earlier sections of the paper. Since this is fairly new concept, this work aims to put forth a proof-of-concep and seeks to answer research questions instead of making 'sotaeesque' comparison. We leave extensive comparisons on other setups and baselines to future work.
Dataset We conduct our experiments on the challenging Natural Questions (NQ) (kwittkwssii et al. 2019) dataset. NQ consists of 307K query- query-document training pairs and 8K validation pairs, where the queries are natural language questions and the documents are Wikipedia articles. Given a question, the retrieval task is to identify the Wikipedia article that answers it. For evaluating how DSI models perform different scales, we construct three sets from NQ form our testbed, namely NQ10K, NQ100K, and NQ320K denoting different numbers of total guery-document pairs in the combined train and validation splits. NQ320K is the full NQ set and uses its predetermined training and validation split for evaluation purposes. Unlike NQ320K, NQ10K and NQ100K constructs randomly sampled validation sets. For all datasets, we use the same docid spacc/budget of 320K tokens for all unstructured atomic and naively structured identifier experiments. Semantically structured identifiers are generated separately for each dataset so as to prevent Icakage of semantic information from larger splits into smaller ones. Text is lowercased. Note that there exists fewer unique documents than query-dar guery-document pairs in these datasets. Please refer to Table 4 (Appendixx which reports the statisties of these datasets.
Metrics We evaluate our models on Hits N where 10} This metric reports the proportion of correct documents ranked in the top N predictions.
Implementation Details All DSI models are initialized using standard pretrained T5 (Raffel et al. 2019) model configurations. The configurations names and corresponding number of model parameters are: Base (0.2B), Large (0.8B), XI. (3B) and XXL. (11B). For unstructured atomic identifiers runs, we initialize the identifiers randomly as new parameters and only finetune the weights implementation for our experiments. The DSI during the indexing stage. We use the Jax/T5X P
models are trained for maximum of IM steps using batch size of 128. We pick the best checkpoint based on retrieval validation performance. Our training hardware consists of 128-256 TPUv4 chips for models above 1B parameters and 128 TPUv3 or TPUv4 chips otherwise. As an estimate, models above 1B parameters typically take about at least a full day for convergence for NQ320K. We tune the learning rate amongst {0.001, 0.0005) and linear warmup amongst {10K. 100K, 200K. 300K) and/or none. Semantically structured identifiers are generated using an 8-layer BERT (Devlin) et al. 2018) model? [3] and the default k-means clustering in scikit-learn. Based on our early ablation experiments of various DSI setting, the main results presented use direct indexing (L 32) and the Inputs2Targes indexing strategy. We present results for all the docid representation methods. Following the main results, we present our ablation studies.
4.1 Baselines
For baselines, we use T5-based dual encoders implemented by (Ni et al. 2021}. We use the gensim package for computing BM25 scores. For the T5-based dual encoders, we train with contrastive learning on the NQ pairs until convergence ( 10K steps) and obtain top-k nearest neighbors with a system similar to ScaNN (Guo et al. 2020). For zero-shot retrieval, we also compare with a state-of- the-art unsupervised baseline, Sentence T5 (Ni et al., (20TT) which have been specially pre-trained with similarity learning task. There two reasons why we consider (Ni et al.) (021)) the relevant dual
encoder baseline for this work rather than other dense retrieval works such as DPR (Karpukhin et al.
https://pppp..org/proooeee
7

---

Table 2: Experimental results NQ document retrieval. DSI outperforms BM25 and Dual Encoder
baselines. Among all the Docid representation methods, Semantic String Docids perform the best.
NQIOK
NQ100K
NQ320K
Model
Size
Params
Method
Hits
Hits
Hits@
Hits@ 10
Hits@ Hits@
BM25
12.4
33.5
20.9
46.4
11.6
34.4
T5
Base
220M
Dual Encoder
16.2
48.6
18.7
65.2
20.5
58.3
T5
Large
800M
Dual Encoder
18.8
55.7
22.3
60.5
22.4
63.3
XI.
3B
Dual Encoder
20.8
59.6
23.3
63.2
23.9
65.8
XXL
Encoder
64.5
24.3
67.3
DSI
Base
250M
Atomic Docid
13.0
38.4
23.8
58.6
20.7
40.9
DSI
Large
800M
Atomic Docid
31.3
59.4
17.1
52.3
11.6
37.6
DSI
XL
Atomic Docid
40.1
76.9
19.0
55.3
61.9
DSI
XXL
IIB
Atomic Docid
39.4
77.0
25.3
67.9
24.0
55.1
DSI
Base
250M
Naive String Docid
28.1
18.7
44.6
21.0
DSI
Large
800M
Naive String Docid
34.7
60.5
21.2
50.7
13.3
33.6
DSI
XL
3B
Naive String Docid
44.7
66. 4
24.0
55.1
16.7
58.1
DSI
XXL
1IB
Naive String Docid
46.7
77.9
27.5
62.4
23.8
55.9
DSI
Base
250M
Semantic String Docid
33.9
57.3
19.0
44.9
27.4
56.6
DSI
Large
800M
Semantic String Docid
37.5
65.1
20.4
50.2
35.6
62.6
DSI
XL
Semantic String Docid
41.9
67.1
22.4
62.2
39.1
60.8
XXL
Semantic String Docid
48.5
72.1
26.9
59.5
70.3
Table 3: Experimental results on Zero-Shot NQ document retrieval. DSI outperforms BM25, T5 embeddings and Sentence the state-of-the-ard for unsupervised similarity modeling, Among Docid representation method, the Atomic Docid performs the best on zero-shot learning.
NQ10K
NQ100K
NQ320K
Model
Size
Method
Hits@ 10
Hits@
@its@
Hits@
Hits@
BM25
12.4
33.5
20.9
46.4
11.6
34.4
T5
XXL
Dual Encoder
0.3
1.3
1.9
8.0
1.1
5.9
SentenceTS
Large
Dual Encoder
17.6
50.7
17.4
50.8
16.9
51.0
DSI
XXL
Atomic Docid
25.7
60.1
23.0
57.3
25.1
56.6
DSI
XXL
Naive String Docid
43.4
67.4
17.4
41.5
9.2
22.6
DSI
XXL
Semantic String Docid
43.9
68.8
11.4
26.6
13.9
31.1
2020). Firstly, we employ the exact identical pretrained model, which allows systematic ablation of the proposed approach without conflating other factors. Scientifically, we believe this comparison against fine-tuned T5 is the best apples to apples comparison that we provide. Secondly, fine-tuned TS dual encoders are considered be architecturally and methodologically very identical DPR (with some minor differences such as parameter sharing but use the same concept of in batcatch negatives).
4.2 Experimental Results
Table reports retrieval results for NQ10K, NQ100K, and NQ320K with finetuning and Table reports zero-shot retrieval results. For zero-shot retrieval, the model is only trained on the indexing task and not the retrieval task, so the model sees no labeled query - docid data points. Section 7f. the Appendix reports extended results regarding the indexing performance and training dynamics of DSI.
Supervised Finetuning Results Our results show that DSI outperforms DE across all dataset sizes. On the small dataset (NQ10K), the performance gap between DSI and DE is large, e.g., the best DSI variant outperforms DE by 2 times. On NQ100K, the gap becomes less prominent with the best DSI model (unstructured atomic identifiers) outperforming DE by +5% Hits @ and Hits @ 10. On the large dataset (NQ320K), the best DSI model (structured semantic identifiers) outperform the best DE model by +66% relative Hits @ 1 and +4.5% Hits @ 10.
Zero-Shot Results Table 3 reports results on zeros-shot retrieval. Recall that zero-shot retrieval is performed by only performing indexing and not the retrieval task. In other words, the model does not any annotated query document pairs. Generally, the best result is obtained by DSI with

---

3x30
Performance (hits@))
Coal Enooder taave))
Model Size (log scale)
Multi-task Ratio
Figure 3: Scaling plots for DSI vs. DE across model sizes. Performance refers to the Hits@ metric.
Figure 4: Effect of multi-task ratio of index- ing to retrieval examples.
Hits@1
Inverted index (L-32)
Set (L=64)
Direct [L=1281
Direct (L-54)
Direct (L=32)
10
15 20
Figure 5: Performance of different document representations. (Referenced in Section 4.2
unstructured atomic identifiers on both NQ100K and NQ320K. The best performance on all NQ datasets outperform well-established unsupervised retrieval baselines such as BM25. Moreover, DSI outperforms unsupervised representation learning methods such as SentenceTSS (Ni et al.) [002]). which is trained to learn similarity-aware representations via contrastive learning. We also note that raw T5 embeddings perform extremely poorly and do not produce reasonable results on the task of unsupervised retrieval. Given that it is generally difficult for unsupervised neural method to outperform BM25, we find these early results very encouraging.
Document Identifiers One key research question in this paper the crucial choice of how to represent docids. Generally, we find that structured semantic identifiers are helpful and improve over unstructured identifiers. When comparing naive versus semantic string identifiers, seems imperative to use semantic identifiers if possible. This is intuitive, since imbuing the target space with semantic structure can facilitate greater ease of optimization and additional unsupervised representation learning methods external knowledge. The competitiveness of unstructured atomic identifiers somewhat mixed and we had some difficulty optimizing such models. We hypothesize that this could possibly be because of the the newly initialized softmax layer and that training such a system from scratch would mitigate these issues. However, we defer this line of investigation to future work. In lieu of the instability and high variance of the unstructured atomic identifiers, the performance is not consistent across the different datasets. Moreover, these docids might also run optimization related quirk. However,
into intermittent non-convergence which trace back we also note that unstructured atomic identifiers perform the best, by a wide margin, on the zero-shot retrieval setup and achieve performance often more than double than that of beam decoding methods.
Indexing Strategies In this section, we explore the effect of different indexing methods (Sec- tion 3.1.1). We run experiments on NQ100K with the different indexing strategies described earlier. Models are trained using the Naive Docid method. Without indexing, the model achieves 0% Hits@ 1. This is intuitive, since Docids are not meaningful without the indexing task. Secondly, the Inputs2Targets and Bidirectional formulation performs the best, with the bidirectional method performing slightly worse (13.5 vs 13.2) compared to the former. Finally, the accuracy with Targets2Inputs and Span Corppution with Docids yield meaningful results (0% accuracy). This

---

goes to show that there can be huge variance across indexing strategies whereby some strategies work reasonably well and some completely do not work at all.
Document Representations In this section, we explore the performance of the different document representation strategies described in Section 3.1.2] FigureBJeports the results on NQ320K. Overall, we find that the direct indexing approach works the best. We also find that is difficult train the inverted index method since the docid is repeatedly exposed to different tokens. We also find that shorter document lengths seem to work well where performance seems substantially dip beyond 64 tokens suggesting that it might be harder to optimize or efficiently memorize when there are larger number of document tokens. Finally, we also find that there was no additional advantage in applying set processing or stopwords preprocessing to the document tokens.
Another interesting insight is how the scaling law of DSI differs from Dual Encoders. Understanding the scaling behaviour of Transformers have garnered significant interest in recent years (Kaplan et al. 2020; Tay et al. 2021 Abnar et al. (021). We find that the gain in retrieval performance obtained from increasing model parameterization DE seems to be relatively small. Conversely, the scaling properties of DSI seems to be more optimistic.
Scaling Laws
Figure 3 plots the scaling behaviour (log scale) of three methods (DE and DSI with naive and semantic IDs). DSI (naive) strongly benefits from scale going from base to XXL and seems to still have headroom for improvement. Meanwhile, DSI (semantic) starts off equally competitive as DE base but performs much better with scale. DE models, unfortunately are more or less plateaued at smaller parameterization.
Interplay Between Indexing and Retrieval Our early experiments showed that first learning the indexing task and then learning the retrieval task in a sequential manner results in mediocre performance. There, we focused on exploring good ratios for co-training the indexing and retrieval tasks together using multi-task learning. Figure I shows the effect of modifying the ratio of indexing to retrieval samples. We find the optimization process is significantly influenced by the interplay between the indexing and retrieval tasks. Setting r too high or low generally resulted in poor performance, We find that rate of 32 generally performed well.
5
Conclusion
This paper proposed the Differentiable Search Index (DSI), a new paradigm for learning an end-to-end search system in unified manner, paving the way for next generation search (Metzler et al.) 10211) We define novel indexing and retrieval tasks that encode the relationship between terms and docids completely within the parameters of Transformer model. The paper proposed a number of different ways to represent documents and docids, and explored different model architectures and model training strategies. Experiments conducted on the Natural Questions data set show that DSI performs favorably against common baselines such as BM25 and dual encoders, both in a standard fine-tuning setup as well as zero-shot setup.
Although the models and results presented here are promising, there a great deal of potential future research that can be explored based this work to improve this approach. For example, it would be interesting to explore alternative strategies for representing documents and docids, as well as to investigate mixture-of-expper models (Du et al., |021; 1 Fedus et al.) 2021 Lepikhin et al.) 2020) for scaling the memory capacity of DSI. One important direction will also be to explore how such models can be updated for dynamic corpora, where documents may be added or removed from the system. Finally it may also be interesting to further investigate DSI as an unsupervised representation learning method and/or memory store for other language models to leverage.
Acknowledgements
The authors would like to thank you Fernando Pereira, Huaixiu Steven Zheng, Sebastian Ruder, Adam D. Lelkes, Ian Wetherbee and Dani Yogatama for their valuable feedback and discussions. We would also like to extend special thanks to Sanket Vaibhav Mehta for additional experimental contributions.
10

---

References
Samira Abnar, Mostafa Dehghani, Behnam Neyshabur, and Hanie Sedghi. Exploring the limits of large scale pre-training. arXiv preprint arXiv:2110.02095, 2021.
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George van den Driessche, Jean- Sann-Bapisste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. arXiv preprint 2 2021.
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. /4165, 2020.
Nicola De Cao, Gautier Izacard, Sebastian Riedel, and Fabio Petroni. Autoregressive entity retrieval. arXiv preprint arXiv:2010.00900 00,0044 2020.
Mostafa Dehghani, Hamed Zamani, Aliaksei Severyn, Jaap Kamps, and W Bruce Croft. Neural ranking models with weak supervision. In Proceedings of the 40th International ACM SIGIR Conference on Research and Development Information Retrieval, pages 65 2017.
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04800. 2018.
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts arXiv preprint arXiv:21 22.22202..690 2021.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961, 2021.
Tianyu Gao, Xingcheng Yao, and Danqi Chen. SimCSE:
Simple contrastive learning of sentence embeddings. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 6894-6910, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653/v1/20211 1en.201.eenpppmaa55555551 URI https
//aclanthology.. . org/ 2021 emnlp-main. 552
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers key-value memories. arXiv preprini arXiv:2012. /4913, 2020.
Daniel Gillick, Alessandro Presta, and Gaurav Singh Tomar. End-to-end retrieval in continuous space. arXiv preprint 1811v000 2018
Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiv Kumar. Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning, 2020. https:tss:///rxii www.oorgg/abbbaaoo /tw./////.... 10396
Kelvin Guu, Kenton Lee, Zora Tung. and Panupong Pasupat. REALM: Retrieval-Augennted
Language Model Pre- Training. In Proceedings of ICML 2020, 2020.
Martin Josifoski, Nicola De Cao, Maxime Peyrard, and Robert West. Genie: Generative information extraction. arXiv preprint arXiv:21 2021.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint (tX.i:200010086111 2020.
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:2004.04490 04490 2020.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin Kenton Lee, Kristina Toutanova, Llion Jones Matthew Kelcey, Ming- Wei Chang, Andrew M Dai, Jakob Uszkoreit, Quoe Le, and Slav Petrov. Natural Questions: a Benchmark for Question Answering Research. In Transactions of the
ACL, 2019.
11

---

Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.
Donald Metzler, Yi Tay, Dara Bahri, and Mare Najork. Rethinking search: making domain experts out of dilettantes. In ACM SIGIR Forum, volume 55, pages 1-27. ACM New York, NY, USA, 2021.
Jianmo Ni, Gustavo Hernández Ábrego, Noah Constant, Ji Ma, Keith B Hall, Daniel Cer, and Yinfei Yang. Sentence-t5: Scalable sentence encoders from pre-trained text-to-text models. arXiv preprint arXiv:2108. 08877, 2021.
Fabio Petroni, Tim Rocktäschel, Patrick Lewis, Anton Bakhtin, Yuxiang Wu, Alexander H Miller, and Sebastian Riedel. Language models as knowledge bases? arXiv preprint 190.0000 2019.
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: 1066,,,,,, 2019.
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, al. Hopfield networks is all you need. arXiv preprint axXiv:2008.02217 2020..
Adam Roberts, Colin Raffel, and Noam Shazeer. How much knowledge can you pack into the parameters of language model? arXiv preprint wxXi:22002.08100 2020.
Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with self- self-supervision for generalization under distribution shifts. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 9229-922444 PMLR, 13-18 Jul 2020. URL itttps://proceedinggs mlr press/ httpsss//vv/////// 1111//nn2001....mmmm wwwssuunnnussssee
Ilya Sutskever, Oriol Vinyals, and Quoe Le. Sequence to sequence learning with neural networks. arXiv preprint arXiv: 1rXhv:140.33255. 2014.
Yi Tay, Mostafa Dehghani, Jinfeng Rao, William Fedus, Samira Abnar, Hyung Won Chung, Sharan Narang, Dani Yogatama, Ashish Vaswani, and Donald Metzler. Scale efficiently: Insights from pre-training and fine-tuning transformers. arXiv preprint arXiv:2109. 10686, 2021.
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng- Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint 2r0iv:2201. 082399, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6008, 2017.
12

---

Checklist
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 descrihe the limitations of your work? [Yes] We highlight that this is proof-
of-concept for new paradigm, and not geared towards SOTA comparisons (Section
Clearly, many open problems still remain such as index updating, scaling to larger datasets,
(c) Did you discuss any potential negative societal impacts of your work? [No] Our work effectively provides the same end functionality existing information retrieval systems.
(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] No theoretical results.
(b) Did you include complete proofs of all theoretical results? [N/A] No thcoretical results
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] All experiments using publicly available T5X codebase, with publicly available Natural Ques- tions dataset. All necessary details were included to reproduce the experiments.
(b) Did you specify all the training details (e.g.. data splits, hyperparameters, how they were chosen)? [Yes] See Section 4 under Implementation Details.
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [No] As these experiments expensive, experiments were not ran multiple times to save compute.
(d) Did you include the total amount of compute and the type of resources used (c.g.. type of GPUs, internal cluster, or cloud provider)? [Yes] Please see Section 4 under Implementation Details.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [No] We comply with the license of all code
repositories and public datasets as detailed in the cited sources.
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A] No new assets introduced at this time.
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A] We did not curate data.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A] We did not curate data.
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "n5tdZDU3-4z", "reviewer": "Reviewer_dTow", "summary": "The authors propose a new information retrieval framework called DSI that encodes the corpus into the parameters of the model and generates the doc-ids for the queries. In addition, this paper studies variation in how documents and their identifiers are represented, variations in training procedures, and the interplay between models and corpus sizes. Experiments show that DSI outperforms strong baselines such as dual encoder models and has strong generalization capabilities.", "questions": "(1)\tIf the corpus is encoded with the model, it becomes difficult to add or delete documents in the corpus. How to solve this problem?\n\n(2)\tFor a large-scale corpus, the parameter size of the model is required to be larger. Will this increase the burden of the online IR system? \n\n(3)\tCan the process of using doc-id to determine the document itself also be understood as a kind of index? How can this process be distinguished from the definition of the index in IR?\n", "limitations": "The limitations of this work have been described in detail. I think how to dynamically update the documents in the corpus and how to deal with large-scale corpus are issues that authors need to consider.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n(1) Propose a novel text-to-text information retrieval paradigm, which has a certain potential for improving information retrieval performance, especially generalization ability.\n\n(2) The experiments are comprehensive and some key variants are discussed.\n\nWeaknesses:\n\n(1) The practical utility of this method needs to be discussed in detail because encoding a large corpus requires a model with large parameters, which makes the model more resource-intensive online. \n\n(2) Using the model to encode the entire corpus makes the interpretability of the IR system worse, and the output becomes more uncontrollable.\n\n(3) In addition, using the model as an index makes it more difficult to add and remove documents from the corpus. For example, a new doc-id is hard to learn by the model if not trained from scratch.", "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": "iQVucrhN_8", "reviewer": "Reviewer_GpFY", "summary": "This paper proposes a sequence-to-sequence style approach to information retrieval, where a transformer-based seq2seq model maps a query (or a document) to a document identifier. The authors consider several approaches to representing both the input documents and the document identifiers, and find that standard document encoding but a “semantically structured” document identifier representation work best. Variants of the proposed model outperform reasonable baselines on Natural Questions-based benchmarks in both fine-tuned and zero shot settings.", "questions": "See the last point above for two questions. In addition:\n- Do you have any insights into what explains the difference in performance between the various settings, both within Table 2 and between Tables 2 and 3?\n- Do you think the semantic string docid approach would work with unseen documents?\n- Extremely minor: Table 4 in the appendix makes it seem like NQ320K has only 228K documents; are these just the documents seen at training time perhaps?", "limitations": "I think the authors may want to address a bit more explicitly how this approach might scale to documents unseen at training time. Otherwise I think limitations are reasonably well addressed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper presents a very interesting and provocative idea.\n- The methods proposed are simple and reasonable (in a good way!).\n- All in all, the paper obtains good performance on the tasks it considers.\n- The paper is written very clearly.\n- The authors present results for a relatively large number of data-size/model-size configurations.\n\nWeaknesses:\n- I think one weakness is that it’s a little bit difficult to understand why the experimental results are what they are. In particular, different model variants (in Table 2) appear to be best in different settings, and the results don’t seem to be monotonic either in the dataset size nor even in the model size in certain cases. Part of the issue here is that the experiments don’t really control for a number of important factors. First, the datasets differ not just in how many total documents they include but also (from Table 4 in the appendix) in the number of training examples. So it’s difficult to ascertain which (if either) is responsible for the change in performance. Similarly, it appears that in the Atomic Docid setting the authors do exact inference, whereas some beam approximation is used for the other two settings. So again it’s difficult to tell if this is a factor too. \n- More minor: I think a natural concern about this approach is what happens with documents that have not been seen at training time, and I’d expect the authors to at least discuss this issue.\n- Finally, and most minor: the authors view training the model to associate a doc id with a document as a form of differentiable indexing. But is there evidence that this can be used for anything? For example, after training, can we feed only a part of the document to the model and reliably retrieve its id? And if not, is it even necessary to train the model for indexing? Perhaps we could just train the model to associate doc ids with *queries*, which is how we use the model at test time.\n\nUpdate after response: thanks for your response; I'm increasing my score. I do continue to think, though, that the paper would be stronger if it addressed the potential confounds of number of documents vs. number of training examples and exact vs. approximate inference.", "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": "dKAr3iRnQN3", "reviewer": "Reviewer_vdLK", "summary": "This paper proposes a sequence-to-sequence architecture that predicts the target document ID directly from the query to replace information retrieval approaches that are based on pipeline retrieve-then-rank strategies. The main idea is to build a differentiable search index where all the information about the documents is stored in the parameters of the transformer language model. To achieve that the authors experiment with different document and document ID representations. Evaluation is performed on corpus sizes of increasing size up to 320K and that it outperforms competitive B25 and dual encoder baselines. Notably the results improve when increasing the model size.", "questions": "Could the authors summarize which design choices performed poorly in the introduction? Specifically, in lines 50-52 and 70-71. The reader needs to go through model section or reach the last section to figure out what did not work. \n\nWhat is the retrieval speed of the proposed system compared to existing alternatives e.g. dense retrieval or BM25? Indexing using BM25 should be really fast and the same holds for dense retrieval, especially with MIPS. The trained model on the other hand should generally be much slower to generate the document ID given large pretrained models are known to be slow during decoding. These potential limitations should be highlighted in the text. \n\nBeyond performance on the retrieval tasks, have the authors looked into what extent the documents are being correctly memorized by the pretrained transformer model? For instance, looking at the predicted document ID for a sample of them when part or all the document is given as input to the decoder. It would be interesting to find out if there are any documents that the mode confuses with each other.\n\nAny insights into why T5 XXL performed that poorly in the zero-shot setting? It is severely under-performing across the board compared to SentenceT5. \n", "limitations": "I'd suggest the authors to discuss the following:\n- More detailed description of limitations compared to dense retrieval and BM25: e.g. in terms of training difficulty, cost of maintaining/hosting, retrieval speed.\n- Potential negative societal impact from memorizing data with pretrained models in an explicit way. Such functionality could indeed have unwanted privacy/confidentiality consequences that deserve some careful thought and discussion. ", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths** \n- The paper is generally well written and proposes a novel differentiable search index paradigm that uses a sequence-to-sequence transformer model to memorize documents by their IDs and then use it to perform information retrieval tasks.\n- Evaluation on datasets of increasing size demonstrate consistent improvements over both dense retrieval and BM25 baselines albeit different variations worked best on different datasets. It is noteworthy that even simple atomic document IDs perform very competitively. \n- In dense information retrieval, it is generally very challenging to beat standard baselines such as BM25  and in some cases retrieval focused pre-training is required on multiple datasets (e.g. [1]); let alone for zero-shot retrieval tasks. \n- The contributions of this paper open up new avenues for research and should be of interest to both machine learning (memorizing and associative memory) and information retrieval communities (new paradigm for dense information retrieval). \n\n**Weaknesses**\n- In terms of experiments, there is no discussion of efficiency limitations of the proposed approach and the ablation studies do not appear to be very thorough; it's unclear how much effort was put into tuning each option. \n\n- Most of the experiments are sufficiently described but there is no code shared in the supplemental material or statement in the paper about code release which could be an issue for reproducibility. \n\n[1] Unsupervised Dense Information Retrieval with Contrastive Learning: https://arxiv.org/pdf/2112.09118.pdf", "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": "VHdIf4A8Znn", "reviewer": "Reviewer_iAXN", "summary": "The paper proposes a generative approach for document retrieval task. Traditional retrieval systems use a dual encoder to map the query and the document to the same vector space, and then perform nearest neighbor search to find the closest document. In contrast to such systems, the proposed system retrieves the target document by mapping each query to the document id in a generative way. In this case, indexing can be done by memorizing the corpus in the parameters of the language model (similar to associative memory). This mechanism has several inherent benefits, one of them being no need to perform indexing which takes a lot of space and time. The paper explores several different ways to index the corpus and represent the document id. The paper shows that inputs2target indexing (learning to generate the target document id with the document content as the input) and semantic document identifier (via clustering) work the best. The results are comparable or better than the dual encoder baselines in Natural Questions 320k dataset (a subset of NQ 320k) and also show promising results in zero-shot setting where the query-document pairs are never observed.", "questions": "Questions\n- L201: Doesn't Semantically Structured Identifiers also require offline indexing? Then I think the argument that DSI does not requiring offline indexing is not entirely true.\n- It would have been easier for me to compare the model with DPR if Table 2 reports Hits@5 and/or Hits@20. Would you be able to provide them?\n\nTypos\n- L29: take -> takes\n- L110: \\citep -> \\citet (there are several other places in the paragraph that the authors make the same mistake)\n- L220: space between \"forcing.\" and \"We\"", "limitations": "The paper mentions that what the paper proposes is more of a novel idea and perspective on retrieval than a SOTA-ish contribution. I agree with this.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The proposed method is novel. It is a new perspective to think of document retrieval task as a doc id generation. If this works well, it has the potential of complimenting or even replacing traditional dual encoder retrieval paradigm. So I think the paper has a high impact on the research community.\n- Zero-shot performance is quite promising. Beating BM25 with an unsupervised model is not easy.\n\nWeaknesses\n- The paper experiments on only one dataset (Natural Questions), and on its subset (320k out of 5 million Wikipedia articles). Given that NQ is a very representative dataset in this field, experimenting on only NQ seems fine, but it is not clear why this was not experimented on the full article set. Even if the performance degrades, I think the paper would have looked better with the full Wikipedia results.\n- The dual encoder baselines seem to be weak. The paper uses Sentence T5 (Ni et al., 2021), which is not originally designed for retrieval task. A stronger and more appropriate baseline would have been DPR or something similar, which uses in-batch negative training objective to leverage the benefit of dual encoder training.\n- The paper does not experimentally compare with a very relevant previous work, AER (Cao et al., 2020). While I agree that DSI is more generally applicable, it would have been great if AER and DSI were comparable in a similar experimental setup to gauge the advantage and disadvantage of doc id generation with respect to title generation.\n\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