Title: Learning to Tokenize for Generative Retrieval
Venue: NeurIPS 2023
Accepted: yes
Keywords: Information Retrieval

Document Retrieval

Generative Retrieval
Abstract: As a new paradigm in information retrieval, generative retrieval directly generates a ranked list of document identifiers (docids) for a given query using generative language models (LMs).
How to assign each document a unique docid (denoted as document tokenization) is a critical problem, because it determines whether the generative retrieval model can precisely retrieve any document by simply decoding its docid.
Most existing methods adopt rule-based tokenization, which is ad-hoc and does not generalize well.
In contrast, in this paper we propose a novel document tokenization learning method, GenRet, which learns to encode the complete document semantics into docids.
GenRet learns to tokenize documents into short discrete representations (i.e., docids) via a discrete auto-encoding approach.
We develop a progressive training scheme to capture the autoregressive nature of docids and diverse clustering techniques to stabilize the training process.
Based on the semantic-embedded docids of any set of documents, the generative retrieval model can learn to generate the most relevant docid only according to the docids' semantic relevance to the queries.
We conduct experiments on the NQ320K, MS MARCO, and BEIR datasets.
GenRet establishes the new state-of-the-art on the NQ320K dataset.
Compared to generative retrieval baselines, GenRet can achieve significant improvements on unseen documents.
Moreover, GenRet can also outperform comparable baselines on MS MARCO and BEIR, demonstrating the method's generalizability.

Learning to Tokenize for Generative Retrieval

Abstract

As a new paradigm in information retrieval, generative retrieval directly generates a ranked list of document identifiers (docids) for a given query using generative language models (LMs). How to assign each document a unique docid (denoted as document tokenization ) is a critical problem, because it determines whether the generative retrieval model can precisely retrieve any document by simply decoding its docid. Most existing methods adopt rule-based tokenization, which is ad-hoc and does not generalize well. In contrast, in this paper we propose a novel document tokenization learning method, GENRET, which learns to encode the complete document semantics into docids. GENRET learns to tokenize documents into short discrete representations (i.e., docids) via a discrete auto-encoding approach. We develop a progressive training scheme to capture the autoregressive nature of docids and diverse clustering techniques to stabilize the training process. Based on the semantic-embedded docids of any set of documents, the generative retrieval model can learn to generate the most relevant docid only according to the docids' semantic relevance to the queries. We conduct experiments on the NQ320K, MS MARCO, and BEIR datasets. GENRET establishes the new state-of-the-art on the NQ320K dataset. Compared to generative retrieval baselines, GENRET can achieve significant improvements on unseen documents. Moreover, GENRET can also outperform comparable baselines on MS MARCO and BEIR, demonstrating the method's generalizability.

1 Introduction

Document retrieval plays an essential role in web search applications and various downstream knowledge-intensive tasks by identifying relevant documents to satisfy users' queries. Recently, generative retrieval has emerged as a new paradigm for document retrieval [1, 5, 37, 41, 46, 47] that directly generates a ranked list of document identifiers (docids) for a given query using generative language models (LMs). Unlike dense retrieval [9, 13, 23, 42], generative retrieval presents an end-to-end solution for document retrieval tasks [37]. It also offers a promising approach to better exploit the capabilities of recent large LMs [1, 41].

As shown in Figure 1, document tokenization aims to tokenize each document in the corpus as a sequence of discrete characters, i.e., docids. Document tokenization plays a crucial role in generative retrieval, as it defines how the document is distributed in the semantic space [37]. And it is still an open problem how to define docids. Most previous generative methods tend to employ rule-based document tokenizers, such as generating titles or URLs [5, 46], or clustering results from off-the-shelf document embeddings [37, 41]. Such rule-based methods are usually ad-hoc and do not generalize well. In particular, the tokenization results potentially perform well on retrieving documents that have been seen during training, but generalize poorly to unlabeled documents [17, 20].

𝑑

32 45 21

Transformer

Query

32 45 21

Document

𝑑

𝑅(𝑑|𝑧)

Reconstructed

Reconstruction

Codebook

docid

𝑧

docid

𝑧

Transformer

Codebook

To address the above problem, we propose GENRET, a document tokenization learning framework that learns to tokenize a document into semantic docids in a discrete auto-encoding scheme. GENRET consists of a shared sequence-to-sequence-based document tokenization model, a generative retrieval model, and a reconstruction model. In the proposed auto-encoding learning scheme, the tokenization model learns to convert documents to discrete docids, which are subsequently utilized by the reconstruction model to reconstruct the original document. The generative retrieval model is trained to generate docids in an autoregressive manner for a given query. The above three models are optimized in an end-to-end fashion to achieve seamless integration.

There are usually two challenges when using auto-encoding to optimize a generative retrieval model: (i) docids with an autoregressive nature, and (ii) docids with diversity. To address the first challenge and also to stabilize the training of GENRET, we devise a progressive training scheme. This training scheme allows for a stable training of the model by fixing optimized prefix docids z <t . To optimize the docids at each step, three proposed losses are utilized: (i) a reconstruction loss for predicting the document using the generated docid, (ii) a commitment loss for committing the docid and avoiding forgetting, and (iii) a retrieval loss for optimizing the retrieval performance end-to-end. To address the second challenge, we propose a parameter initialization strategy and a re-assignment of the docid based on a diverse clustering technique to increase the diversity of the generated docids.

We conduct extensive experiments on three well-known document retrieval benchmark datasets, NQ320K [15, 37], MS MARCO [4, 46], and BEIR [38]. GENRET attains superior retrieval performance against state-of-the-art generative retrieval models on NQ320K. GENRET achieves +14% relative improvements on the unseen test set of NQ320K over the best generative retrieval baseline. Experiments on MS MARCO and six BEIR datasets also show that GENRET outperforms existing generative methods and achieves competitive results compared to popular dense retrieval models. Experiments on retrieving new documents, analytical experiments, and an efficiency analysis confirm the effectiveness of the proposed model.

We summarize our contributions as follows: (i) We propose GENRET, a generative retrieval model that represents documents as discrete semantic docids. To the best of our knowledge, this is the first tokenization learning method for document retrieval. (ii) We propose an auto-encoding approach, where the docids generated by our tokenization model are reconstructed by a reconstruction model to ensure the docids capture the semantic information of the document. (iii) We devise a progressive training scheme to model the autoregressive nature of docids and stabilize the training process. (iv) Experimental results demonstrate that GENRET achieves significant improvements, especially on unseen documents, over generative retrieval baselines. 2

2 Preliminaries

The document retrieval task can be formalized as the process of retrieving a relevant document d for a search query q from a collection of documents D . Each document d ∈ D is assumed to be a plain text consisting of a sequence of tokens, denoted as d = { d 1 , . . . , d | d | } , where | d | is the total number of tokens in the document. For generative retrieval models, it is usually challenging and computationally

inefficient to directly generate original documents of typically long length. Therefore, most existing approaches rely on a technique named document tokenization , which represents a document d = { d 1 , . . . , d | d | } as a shorter sequence of discrete tokens (docid) z = { z 1 , . . . , z t , . . . , z M } , where each token z t is as a K -way categorical variable, with z t ∈ [1 , 2 , . . . , K ] , and M is the length of the docid. See Figure 1 for an example of document tokenization with M = 3 and K = 64 .

As an alternative sequence of the original document, the tokenized docid z should satisfy the following two properties: (i) different documents have short but different docids; and (ii) docids capture the semantics of their associated documents as much as possible [37]. Because z is a sequence of a fixed length and usually shorter than the original document d , the model's training and inference can be simplified and more efficient. This paper employs a tokenization model Q : d → z to map d to docid z . More details about Q are provided in Section 3.1. After tokenizing each document to docid z , a generative retrieval model P : q → z learns to retrieve relevant documents by generating a query q to a docid z autoregressively [37].

3 Method

Conventionally, document tokenization is done by a fixed pre-processing step, such as using the title of a document or the results of hierarchical clustering obtained from BERT [5, 37]. However, it has been observed that such ad-hoc document tokenization methods often fail to capture the complete semantics of a document. For example, the title of a web page often does not exist or has low relevance to the content of the web page, and the use of clustering-based docids arbitrarily defines the document in discrete space.

In this paper, we propose GENRET, a novel tokenization learning method based on discrete autoencoding, to learn semantic docid in a fully end-to-end manner. Figure 1 gives an overview of the proposed method. The proposed GENRET comprises three main components: (i) a sequence-tosequence based retrieval model P ( z | q ) , (ii) a document tokenization model Q ( z | d ) , and (iii) a reconstruction model R ( d | z ) . The document tokenization model tokenizes a document d into unique discrete variables z , and the retrieval model is trained to generate the latent variables z for a given query q . In addition, the reconstruction model is used to re-generate the original document from z to ensure z captures the semantics of the original document as much as possible.

3.1 Model architecture

Following DSI [37], we employ an encoder-decoder Transformer to implement the generative retrieval model. Specifically, given an input text d 3 , the T5-based tokenization model encodes d and a prefix of docid z <t and continuously produces latent representation d t of d at time step t :

where D denotes the hidden size of the model. Encoder( d ) denotes the output of the Encoder .

Then, the tokenization model generates a token for each document based on d t . At each timestep t , we define an external embedding matrix named codebook E t ∈ R K × D , where K is the size of the discrete latent space. There are K embedding vectors e t,j ∈ R D , j ∈ [ K ] , and each vector e t,j can be regarded as the centroid of a segmentation.

Based on the codebook E t , the discrete latent variable z t at timestep t is calculated by a dot-product look-up using the codebook E t :

where Q ( z t = j | z <t , d ) denotes the probability of tokenizing d to a particular value j ∈ [ K ] at timestep t , Softmax j is a softmax function to output the probability of axis j .

Document reconstruction model. The docid generated by the tokenization model Q is required to capture the semantic information of the document. To this end, we propose an auto-encoding training scheme, where a reconstruction model R : z → d that predicts d using z is designed to force the tokenization model Q : d → z to reproduce a docid z that can be reconstructed back-to-the original document.

The input of the reconstruction model is docid z , and the output is its associated document d . We first embed z into representation matrix z = { z 1 , . . . , z M } ∈ R M × D using the codebook of the tokenization model:

where each t ∈ [ M ] , z t = e t,z t ∈ R D is the embedding vector of z t in the t -step codebook E t .

Wethen devise a retrieval-based reconstruction model that predicts the target document d by retrieving it from document collection D , based on the inputs z . The relevance score between the input docid z and the target document d is defined as follows:

where S ( z <t ) is a sub-collection of D consisting of documents that have a docid prefix that is the same as z <t . d ∗ ∈ S ( z <t ) represents a document from the sub-collection S ( z <t ) . d t and d ∗ t are continuous representations of documents d and d ∗ , respectively, as defined in Eq. 1. The operator sg( · ) is the stop gradient operator to prevent gradient back propagation. Intuitively, R ( d | z ) is designed to retrieve a specific document d from a set of documents S ( z <t ) at each timestep t . The set S ( z <t ) only includes those documents that are assigned the same docid prefix z <t as the target document d . By utilizing this loss function, at each step t , the model is facilitated to learn the residual semantics of the documents not captured by the previous docid z <t .

3.2 Model optimization

For the document tokenization model Q ( z | d ) , generative retrieval model P ( z | q ) , and reconstruction model R ( d | z ) , jointly optimizing these three models using auto-encoding is challenging due to the following reasons: (i) Learning docids in an autoregressive fashion . On one hand, the prediction of the z t at time t relies on previously predicted docids z <t , which is often under-optimized at the beginning and rapidly changes during training, making it difficult to reach convergence. On the other hand, simultaneously optimizing z makes it challenging to guarantee a unique docid assignment. Hence, to stabilize the training of GENRET, we devise a progressive training scheme (see Section 3.2.1). (ii) Generating docids with diversity . Optimizing the model using auto-encoding often leads to unbalanced docid assignment: a few major docids are assigned to a large number of documents and most other docids are rarely assigned. Such a sub-optimal distribution of docids affects the model distinguishability, which in turns triggers length increments of docids in order to distinguish conflicting documents. We introduce two diverse clustering techniques to ensure docid diversity (see Section 3.2.2).

3.2.1 Progressive training scheme

To optimize each of the three models listed above in an autoregressive manner, we propose a progressive autoencoding learning scheme, as illustrated in Figure 2. The whole learning scheme contains M learning steps with respect to the final docid in M -token. And the docid z T at step T ∈ [ M ] is learned and optimized at the corresponding learning step. Besides, at each step T ∈ [ M ] , the docid z T and the model parameters associated with z T generation are updated, while previously produced docids z <T and other parameters are kept fixed. By progressively performing the above process, we can finally optimize and learn our models.

At each optimization step, say the T -step, we devise the learning objective for document tokenization consisting of three loss functions detailed below.

Reconstruction loss. We utilize the reconstruction model R ( d | z ) as an auxiliary model to learn to optimize the docid generation, whose main goal is capturing as much semantics in the docid as

possible. Therefore, we define a reconstruction loss function of step T as follows:

where ˆ z ≤ T is the first T representations of the z , and only the variable z T is optimized in step T . Q ( z t = j | z <t , d ) is defined in Eq. 2. And the document tokenization model Q can therefore be optimized when minimizing L Rec.

Of note, since the computation involves a non-differentiable operation ( arg max( · ) ), we apply straight-through gradient estimation to back-propagate the gradient from reconstruction loss to d T following [39, 44], which copies the gradient of z T directly to d T . Specifically, the gradients to document representation d T are defined as ∂ L Rec ∂ d T := ∂ L Rec ∂ z T . And the gradients to the codebook embedding e T,j are defined as ∂ L Rec ∂ e T,j := 1 z T = j ∂ L Rec ∂ z T .

Commitment loss. In addition, to make sure the predicted docid commits to an embedding and to avoid models forgetting previous docid z <t , we add a commitment loss as follows:

Retrieval loss. For the generative retrieval model P , we jointly learn it together with the document tokenization model Q , where P learns to generate the docids of relevant documents d given a query q . Specifically, suppose ( q, d ) are a query and relevant document pair; we define the learning objective of retrieval model P as:

where the first term is a ranking-oriented loss enhancing the model using ( q, d ) pair; d ∗ is an in-batch document sampled from the same training mini-batch B ; q T and d T denote the representation of q and d at timestep T . The second term is the cross-entropy loss for generating docid z based on q .

The final loss we use at stepT is the sum of reconstruction loss, commitment loss, and retrieval loss:

3.2.2 Diverse clustering techniques

To ensure diversity of generated docids, we adopt two diverse clustering techniques-codebook initialization and docid re-assignment at each progressive training step, where codebook initialization mainly aims to increase the balance of semantic space segmentation, and the docid re-assignment mainly aims to increase the balance of docid assignments.

Codebook initialization. In order to initialize the codebook for our model, we first warm-up the model by passing the continuous representation d T to the reconstruction model instead of the docid representation z T as defined in Eq. 3. During this warm-up phase, we optimize the model using the reconstruction loss L Rec and commitment loss L Com. Next, we collect the continuous representations d T of all documents in D , and cluster them into K groups. The centroids of these clusters are then used as the initialized codebook E T . To balance the initialized docid distribution, we utilize a diverse constrained clustering algorithm, Constrained K-Means , which first normalizes the embeddings of each prefix group, and modifies the cluster assignment step (E in EM) by formulating it as a minimum cost flow (MCF) linear network optimization problem [2].

Docid re-assignment. In order to assign docids to a batch of documents, we modify the dot-product look-up results in Eq. 2 by ensuring that the docid for different documents in the batch are distinct following the method described in [6, 44]. Specifically, let D t = { d (1) t , . . . , d ( B ) t } ∈ R B × D denote the continuous representation of a batch of documents with batch size of B . The dot-product results are represented by H = D t · E ⊤ t ∈ R B × K . To obtain distinct docids, we calculate an alternative H ∗ = Diag( u ) exp( H ϵ ) Diag( v ) , where u and v are re-normalization vectors in R K and R B , respectively. The re-normalization vectors are computed via the iterative Sinkhorn-Knopp algorithm [8]. Finally, H ∗ is used instead of H in the Softmax and arg max (Eq. 2) operations to obtain the docid z t .

4 Experimental Setup

4.1 Datasets and evalutaion metrics

We conduct experiments on three well-known document retrieval datasets: NQ320K [15], MS MARCO [4], and BEIR [38]. We divide the test set of NQ320K into seen test and unseen test , based on whether the target documents of the query have annotated queries in the training data, to test the generalization ability of the model on unlabeled documents. More details about data pre-processing and data statistics are listed in the Appendix A.

On NQ320K, we use Recall@{1,10,100} and Mean Reciprocal Rank (MRR)@100 as evaluation metrics, following [41]. On MS MARCO, we use Recall@{1, 10, 100} and MRR@10 as evaluation metrics, following [46]. On BEIR, we use nDCG@10 as the main metrics and calculate the average nDCG@10 values across multiple downstream sub-datasets as overall metrics.

4.2 Baselines

Weconsider three types of baselines: sparse retrieval methods, dense retrieval methods, and generative retrieval methods. The sparse retrieval baselines are: BM25 [32] and DocT5Query [24]. The dense retrieval baselines are: DPR [13], ANCE [42], Sentence-T5 [22], GTR [23], and Contriever [11]. The generative retrieval baselines are: GENRE [5], DSI [37], SEAL [1], CGR-Contra [17], DSIQG [47], NCI [41], and Ultron [46]. The following three baselines use the same pre-trained LM T5 as GENRET: (i) Sentence-T5 outputs continuous vectors, (ii) GENRE outputs document titles, and (iii) DSI-QG outputs clustering IDs, while GENRET outputs docids learned using the proposed tokenization method. See Appendix B for more details on the other baselines.

4.3 Implementation details

Hyper-parameters. In our experiments, we utilize the T5-Base model [27] as the base Transformer and initialize a new codebook embedding E t for each time step. The parameters of both the encoderdecoder and codebook are shared between the tokenization model and the retrieval model. We set the number of clusters to be K = 512 for all datasets, with the length of the docid M being dependent on the number of documents present. For datasets containing a larger number of candidate documents, a larger value of M is set to ensure that all documents are assigned unique document ids. In the docid re-assignment, the hyper-parameter ϵ is set to 1 . 0 , and the Sinkhorn-Knopp algorithm is executed for 100 iterations.

Indexing with query generation. Following previous work [47, 41, 40], we use query generation models to generate synthetic (query, document) pairs for data augmentation. Specifically, we use the pre-trained query generation model from DocT5Query [24] to augment the NQ and MS MARCO datasets. In query generation, we use nucleus sampling with parameters p = 0 . 8 , t = 0 . 8 and generate five queries for each document in the collection. For the BEIR datasets, we use the queries generated by GPL [40]. GPL uses a DocT5Query [24] generator trained on MS MARCO to generate about 250K queries for each BEIR dataset. Note that the query generator used for BEIR is purely trained on MS MARCO (without using any training data of BEIR) and thus conforms to the zero-shot setting of BEIR [38, 40].

Training and inference. The proposed models and the reproduced baselines are implemented with PyTorch 1.7.1 and HuggingFace transformers 4.22.2. We optimize the model using AdamW and set the learning rate to 5 e -4 . The batch size is 256 , and the model is optimized for up to 500k steps for each timestep. During training, we pre-gather documents which share the same docid prefix into a batch. Therefore, the reassignment strategy is applied to a batch, where we aim to have documents with as diverse IDs as possible. We add a factor of 0.1 to the reconstruction losses to balance the scale. In progressive training, we first warm up the model for 5K steps and then initialize the codebook using the clustering centroids as mentioned in Section 3.2.1. We use constrained clustering 4 to obtain diverse clustering results. During inference, we use beam search with constrained decoding [5] and a beam size of 100 .

5 Experimental results

5.1 Main results

Results on NQ320K. In Table 1, we list the results on NQ320K. GENRET outperforms both the strong pre-trained dense retrieval model, GTR, and the previous best generative retrieval method, NCI, thereby establishing a new state-of-the-art on the NQ320K dataset. Furthermore, our results reveal that existing generative retrieval methods perform well on the seen test but lag behind dense retrieval methods on the unseen test. For example, NCI obtains an MRR@100 of 76 . 8 on the seen test, which is higher than the MRR@100 of 65 . 3 obtained by GTR-Base. However, on unseen test data, NCI performs worse than GTR-Base. In contrast, GENRET performs well on both seen and unseen test data. This result highlights the ability of GENRET to combine the advantages of both dense and generative retrieval by learning discrete docids with semantics through end-to-end optimization.

Results on MS MARCO. Table 2 presents the results on the MS MARCO dataset. GENRET outperforms generative retrieval methods such as Ultron and dense retrieval baselines such as ANCE and Sentence-T5. Furthermore, previous generative retrieval methods (e.g., GENRE, Ultron) utilizing metadata such as the title and URL, while exhibiting decent performance on the NQ320K dataset, underperform in comparison to dense retrieval and sparse retrieval methods on the MS MARCO dataset. This may be because the NQ320K dataset retrieves Wikipedia documents, where metadata like the title effectively capture the semantics of the document. In the case of the MS MARCO dataset, which is a web search dataset, the metadata often does not adequately characterize the documents, resulting in a decline in performance of the generative retrieval model. In contrast, GENRET learns to generate semantic docids that effectively enhance the generative retrieval model.

Results on BEIR. Table 3 lists the results of the baselines and GENRET on six datasets of BEIR. These datasets represent a diverse range of information retrieval scenarios. On average, GENRET outperforms strong baselines including BM25 and ST5 GPL, and achieves competitive results compared to previous-best sparse and dense retrieval methods. Additionally, GENRET demonstrates a significant improvement over the previous generative retrieval model GENRE that utilizes titles as docids. Furthermore, GENRE performs poorly on some datasets, such as BEIR-Covid and BEIRSciDocs. This may be because the titles of the documents in these datasets do not adequately capture their semantic content.

5.2 Analytical experiments

We further conduct analytical experiments to study the effectiveness of the proposed method.

In Figure 3 (left), we plot the frequencies of docids at the first timestep of various learning methods. We label each method using a box with a docid and a diversity metric d , which is calculated by: d = 1 -1 2 n ∑ K j =1 | n j -n u | , where |·| represents the absolute value, n denotes the total number of documents, n j denotes the number of documents that have a docid = j , and n u = n K is the expected number of documents per docid under the uniform distribution.

The results demonstrate the superiority of GENRET (represented by the yellow line) in terms of distribution uniformity. It uses all the potential docid k = 512 and achieves the highest diversity metric with a value of d = 0 . 90 . The method without docid reassignment also yields a relatively balanced distribution, with a diversity metric of d = 0 . 77 . However, the distribution of the method without diverse codebook initialization is highly uneven, which could be due to the fact that most of the randomly initialized codebook embeddings are not selected by the model during the initial training phase, leading to a lack of update and further selection in subsequent training. Additionally, the models without diverse clustering tend to converge to a trivial solution where all documents are assigned the same docid.

In Figure 3 (right), the results of two ablated variants are presented. First, GENRET w/o learning is a generative model that has been trained directly using the final output docid from GENRET, without utilizing the proposed learning scheme. Its retrieval performance is comparable to that of GENRET on seen test data; however, it is significantly lower on unseen test data. This variant demonstrates that the generative retrieval model jointly trained with auto-encoding objectives can represent documents more sensibly based on semantics. This could enhance performance on the less-optimized documents. Contrarily, the parameters obtained via cross-entropy loss on docid generation tasks are less effective in conveying the semantic information of documents. Secondly, GENRET w/ T5-Small uses a small model, and its performance is inferior to that of GENRET using T5-Base. However, the gap between the performance on seen and unseen test data is smaller, which could be attributed to the limited fitting capacity of the small model.

In Appendix C, we evaluate the proposed model's capacity to retrieve new documents and find that it performs well in adapting to new documents compared to existing document tokenization approaches. Additionally, in Appendix D, we analyze the efficiency of various retrieval models in comparison to the various baselines in terms of memory usage, offline, and online latency, and show the advantages of the proposed model.

Table 3: nDCG@10 results on BEIR. The results of the methods marked with † are from our own re-implementation, other results are cited from the original paper or implemented using official code. ST5 GPL denotes Sentence-T5 trained on GPL datasets [40].

5%

4%

3%

2%

1%

0%

0

100

200

300

400

500

GenRet

GenRet w/o Reassignment

GenRet w/o Initialization

Docid

Frequency

# = 494

𝒹

= 0.77

# = 34

𝒹

= 0.07

# = 1

𝒹

= 0.00

# = 512

𝒹

= 0.90

75

65

55

45

Recall@1

GenRet

GenRet

w/o learning

Full test

Seen test

Unseen test

GenRet

5.3 Qualitative analysis

Figure 4 (left) illustrates the document content (title) and the corresponding docid generated by GENRET on the NQ320K dataset. We observe that documents with more similar docids tend to have more relevant content. For example, documents with docids starting with 338-173 are related to Email, such as Email marketing , Mail merge , and Email address , while documents with docids starting with 338 relate to information exchange methods (e.g., Business letter , US Postal Service , and Postage stamps ), representing a more generalized semantics than Email alone.

Figure 4 (right) illustrates a word cloud of documents grouped by docid prefixes. It is evident that documents within the same group are semantically related. For example, major words for documents with the docid prefix 338 are mail and stamp . When a second-level docid 173 is added, the corresponding documents become more specifically related to email . With the addition of a third-level docid 1 , the document group becomes specifically associated with email marketing (docid: 338-173-1 ). Similar patterns can be observed in the other three cases. The case study shows that there is a hierarchical semantic structure within the learned docids.

Figure 5 in Appendix E visualizes the codebook embedding and document embedding. We see that the codebook embedding appears to distribute uniformly within the document representation space, producing meaningful clusters when documents are categorized by docids. We also find that the uniformity of the embedding distribution decreases with increasing docid-length. This may be due to the fact that uniform segmentation is more challenging as the semantic granularity becomes finer.

0

338

173

1

448

508

379 - 392

502 - 222

224 - 106

Mail merge

Email address

Business letter

US Postal Service

Postage stamps

196 - 140 - 409

293 - 433 - 204

Human fertilization

Document Title

dueps

addressur

email

currentusing

used.

mail

marToknge

iddresses

hemail

customer

sending

60

市

Possoletter

EunitedStates

may

Docid:338

Docid:338-173

Docid: 338 -173 - 1

U

Externalslinkdit

New York

anye

Presidentiru

Contents hide

United States

nceDirector

United States

Education!

apng

Contentshide

References External

useform

serierie

name

United States

Eaterial

-office

Docid: 294

Martoent

MitER

servicenOffice

Docid:294-502

Departmentpregr

Sodityfederal

Seeretery

S

Department

Docid:294-502-391

historyehene

DoEDOrganizati

HealthsplitCabin

Docid: 0

teamrst

firstr

Carolina Panthers

season season

present

NFC

Football Teague

NationalFootball

season

divis1on

NFLS-team

AEC

Charlotte"

NFCLeague

South

zPantherhio

Football

Carolina-

Docid:254

Docid:254-413

Docid:254-413-465

6 Related work

Sparse and dense retrieval. Traditional sparse retrieval calculates the document score using term matching metrics such as TF-IDF [31], query likelihood [16], or BM25 [32]. Sparse retrieval is widely used in practice due to its efficiency, but often suffers from the lexical mismatches [18]. Dense

retrieval (DR) addresses this by presenting queries and documents in dense vectors and calculating their similarities with the inner product or cosine similarity [13]. Various techniques have been proposed to improve DR models, such as hard negative mining [42, 26], late interaction [14, 34], knowledge distillation [10, 19], and pre-training [30, 23, 11]. Despite their success, DR approaches have several limitations [5, 21]: (i) DR models employ an index-retrieval pipeline with a fixed search procedure (MIPS), making it difficult to optimize the model end-to-end [37, 41]. (ii) Training DR models relies on contrastive learning [13] to distinguish positives from negatives, which is inconsistent with large LMs training objectives [3] and fails to fully utilize the capabilities of pretrained LMs [1, 35].

Generative retrieval. Generative retrieval is gaining attention. It retrieves documents by generating their docid using a generative model like T5. Generative retrieval presents an end-to-end solution for document retrieval tasks [37, 21] and allows for better exploitation of the capabilities of large generative LMs [1]. Cao et al. [5] first propose an autoregressive entity retrieval model to retrieve documents by generating titles. Tay et al. [37] propose a differentiable search index (DSI) and represent the document as atomic id, naive string, or semantic string. Bevilacqua et al. [1] suggest using arbitrary spans of a document as docids. Additionally, multiple-stage pre-training [7, 46], query generation [41, 47, 46], contextualized embedding [17], and continual learning [20], have been explored in recent studies. Recently, Tang et al. [36] introduce a query-based docid and the rehearsal-based document indexing to improve DSI. However, existing generative retrieval models have a limitation in that they rely on fixed document tokenization to produce docids, which often fails to capture the semantic information of a document [37]. It is an open question how one should define the docids. To further capture document semantics in the docid, we propose document tokenization learning methods. The semantic docid is generated by the proposed discrete auto-encoding learning scheme in an end-to-end manner. Concurrently, Rajput et al. [28] propose a RQ-VAE module to produce semantic docids for generative recommender systems. As a comparison, our proposed method jointly models the tokenization and retrieval tasks with shared parameters to better align the model's representation of the two tasks and are optimized in an end-to-end manner.

Discrete representation learning. Learning discrete representations using neural networks is an important research area in machine learning. For images, Rolfe [33] proposes the discrete variational autoencoder, and VQ-VAE [39] learns quantized representations via vector quantization. DallE [29] uses an autoregressive model to generate discrete image representation for text-to-image generation. Recently, representation learning has attracted considerable attention in NLP tasks, for tasks such as machine translation [48], dialogue generation [45], and text classification [12, 43]. For document retrieval, RepCONC [44] uses a discrete representation learning method based on constrained clustering for vector compression. We propose a document tokenization learning method for generative retrieval, which captures the autoregressive nature of docids by progressive training and enhances the diversity of docids by diverse clustering techniques.

7 Conclusions

This paper has proposed a document tokenization learning method for generative retrieval, named GENRET. The proposed method learns to tokenize documents into short discrete representations (i.e., docids) via a discrete auto-encoding approach, which ensures the semantics of the generated docids. A progressive training method and two diverse clustering techniques have been proposed to enhance the model's training. Empirical results on various document retrieval datasets have demonstrated the effectiveness of the proposed method. Especially, GENRET achieves outperformance on unseen documents and can be well generalized to multiple retrieval tasks.

The limitations of this work include experiments only on moderately sized datasets like NQ320K. The employed data are in sufficient quantity to validate the effectiveness of the proposed method, but application to larger-scale data may require more model parameters and computational resources. Another aspect for improvement is the generalization of the model to unoptimized document collections in different types or domains, as compared to continuous embedding approaches. We recognize this as an important research question for generative retrieval but believe that addressing this is beyond the scope of this paper. In future work, we would like to extend the approach to large document collections. We also plan to explore generative pre-training for document tokenization using large-scale language models. Additionally, we intend to investigate the dynamic adaptation of docid prefixes for progressive training.

References

Michele Bevilacqua, Giuseppe Ottaviano, Patrick Lewis, Wen tau Yih, Sebastian Riedel, and Fabio Petroni. Autoregressive search engines: Generating substrings as document identifiers. In NeurIPS 2022 , 2022.

Paul S. Bradley, Kristin P. Bennett, and Ayhan Demiriz. Constrained k-means clustering. Microsoft Research , 2000.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, T. J. Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In NeurIPS 2020 , 2020.

Daniel Fernando Campos, Tri Nguyen, Mir Rosenberg, Xia Song, Jianfeng Gao, Saurabh Tiwary, Rangan Majumder, Li Deng, and Bhaskar Mitra. MS MARCO: A Human Generated MAchine Reading COmprehension Dataset. ArXiv , abs/1611.09268, 2016.

Nicola De Cao, Gautier Izacard, Sebastian Riedel, and Fabio Petroni. Autoregressive entity retrieval. In ICLR 2021 , 2021.

Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. ArXiv , abs/2006.09882, 2020.

Jiangui Chen, Ruqing Zhang, Jiafeng Guo, Yiqun Liu, Yixing Fan, and Xueqi Cheng. CorpusBrain: Pre-train a generative retrieval model for knowledge-intensive language tasks. In CIKM 2022 , 2022.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In NIPS 2013 , 2013.

Daniel Gillick, Alessandro Presta, and Gaurav Singh Tomar. End-to-end retrieval in continuous space. ArXiv , abs/1811.08008, 2018.

Sebastian Hofstätter, Sheng-Chieh Lin, Jheng-Hong Yang, Jimmy J. Lin, and Allan Hanbury. Efficiently teaching an effective dense retriever with balanced topic aware sampling. In SIGIR 2021 , 2021.

Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Unsupervised dense information retrieval with contrastive learning. In TMLR , 2022.

Shuning Jin, Sam Wiseman, Karl Stratos, and Karen Livescu. Discrete latent variable representations for low-resource text classification. In ACL 2020 , 2020.

Vladimir Karpukhin, Barlas O˘ guz, Sewon Min, Patrick Lewis, Ledell Yu Wu, Sergey Edunov, Danqi Chen, and Wen tau Yih. Dense passage retrieval for open-domain question answering. In EMNLP 2020 , 2020.

Omar Khattab and Matei Zaharia. ColBERT: Efficient and effective passage search via contextualized late interaction over BERT. In SIGIR 2020 , 2020.

Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur P. 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, Quoc V. Le, and Slav Petrov. Natural questions: A benchmark for question answering research. TACL , 7:453-466, 2019.

John Lafferty and ChengXiang Zhai. Document language models, query models, and risk minimization for information retrieval. In SIGIR 2001 , 2001.

Hyunji Lee, Jaeyoung Kim, Hoyeon Chang, Hanseok Oh, Sohee Yang, Vladimir Karpukhin, Yi Lu, and Minjoon Seo. Nonparametric decoding for generative retrieval. In Findings of ACL 2023 , 2023.

Jimmy Lin, Rodrigo Nogueira, and Andrew Yates. Pretrained Transformers for Text Ranking: BERT and Beyond . Springer Nature, 2022.

Yuxiang Lu, Yiding Liu, Jiaxiang Liu, Yunsheng Shi, Zhengjie Huang, Shi Feng, Yu Sun, Hao Tian, Hua Wu, Shuaiqiang Wang, Dawei Yin, and Haifeng Wang. ERNIE-Search: Bridging cross-encoder with dual-encoder via self on-the-fly distillation for dense passage retrieval. ArXiv , abs/2205.09153, 2022.

Sanket Vaibhav Mehta, Jai Gupta, Yi Tay, Mostafa Dehghani, Vinh Quang Tran, Jinfeng Rao, Marc-Alexander Najork, Emma Strubell, and Donald Metzler. DSI++: Updating transformer memory with new documents. ArXiv , abs/2212.09744, 2022.

Donald Metzler, Yi Tay, and Dara Bahri. Rethinking search: Making domain experts out of dilettantes. In SIGIR 2021 , 2021.

Jianmo Ni, Gustavo Hernández Ábrego, Noah Constant, Ji Ma, Keith B. Hall, Daniel Matthew Cer, and Yinfei Yang. Sentence-T5: Scalable sentence encoders from pre-trained text-to-text models. In Findings of ACL 2022 , 2022.

Jianmo Ni, Chen Qu, Jing Lu, Zhuyun Dai, Gustavo Hernández Ábrego, Ji Ma, Vincent Zhao, Yi Luan, Keith B. Hall, Ming-Wei Chang, and Yinfei Yang. Large dual encoders are generalizable retrievers. In EMNLP , 2022.

Rodrigo Nogueira1 and Jimmy Lin. From doc2query to doctttttquery. https://cs. uwaterloo.ca/~jimmylin/publications/Nogueira_Lin_2019_docTTTTTquery-v2. pdf , 2019.

Fabio Petroni, Aleksandra Piktus, Angela Fan, Patrick Lewis, Majid Yazdani, Nicola De Cao, James Thorne, Yacine Jernite, Vassilis Plachouras, Tim Rocktaschel, and Sebastian Riedel. KILT: A benchmark for knowledge intensive language tasks. In NAACL 2021 , 2021.

Yingqi Qu, Yuchen Ding, Jing Liu, Kai Liu, Ruiyang Ren, Xin Zhao, Daxiang Dong, Hua Wu, and Haifeng Wang. RocketQA: An optimized training approach to dense passage retrieval for open-domain question answering. In NAACL 2021 , 2021.

Colin Raffel, Noam M. Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. JMLR , 2020.

Shashank Rajput, Nikhil Mehta, Anima Singh, Raghunandan H. Keshavan, Trung Hieu Vu, Lukasz Heldt, Lichan Hong, Yi Tay, Vinh Q. Tran, Jonah Samost, Maciej Kula, Ed H. Chi, and Maheswaran Sathiamoorthy. Recommender systems with generative retrieval. In NeurIPS 2023 , 2023.

Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In ICML 2020 , 2020.

Nils Reimers and Iryna Gurevych. Sentence-bert: Sentence embeddings using siamese bertnetworks. In EMNLP 2019 , 2019.

Stephen E. Robertson and Steve Walker. On relevance weights with little relevance information. In SIGIR 1997 , 1997.

Stephen E. Robertson and Hugo Zaragoza. The probabilistic relevance framework: BM25 and beyond. Found. Trends Inf. Retr. , 2009.

Jason Tyler Rolfe. Discrete variational autoencoders. In ICLR 2017 , 2017.

Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. ColBERTv2: Effective and efficient retrieval via lightweight late interaction. In NAACL 2022 , 2022.

Weiwei Sun, Lingyong Yan, Xinyu Ma, Shuaiqiang Wang, Pengjie Ren, Zhumin Chen, Dawei Yin, and Zhaochun Ren. Is ChatGPT good at search? Investigating large language models as re-ranking agents. In EMNLP 2023 , 2023.

Yubao Tang, Ruqing Zhang, Jiangui Guo, Jiangui Chen, Zuowei Zhu, Shuaiqiang Wang, Dawei Yin, and Xueqi Cheng. Semantic-enhanced differentiable search index inspired by learning strategies. In SIGIR 2023 , 2023.

Yi Tay, Vinh Quang Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Gupta, Tal Schuster, William W. Cohen, and Donald Metzler. Transformer memory as a differentiable search index. In NeurIPS 2022 , 2022.

Nandan Thakur, Nils Reimers, Andreas Ruckl'e, Abhishek Srivastava, and Iryna Gurevych. Beir: A heterogenous benchmark for zero-shot evaluation of information retrieval models. In NeurIPS 2021 , 2021.

Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In NIPS 2017 , 2017.

Kexin Wang, Nandan Thakur, Nils Reimers, and Iryna Gurevych. GPL: Generative pseudo labeling for unsupervised domain adaptation of dense retrieval. In NAACL 2022 , 2022.

Yujing Wang, Ying Hou, Hong 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, and Mao Yang. A neural corpus indexer for document retrieval. In NeurIPS 2022 , 2022.

Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul Bennett, Junaid Ahmed, and Arnold Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval. In ICLR 2021 , 2021.

Erxin Yu, Lan Du, Yuan Jin, Zhepei Wei, and Yi Chang. Learning semantic textual similarity via topic-informed discrete latent variables. In EMNLP 2022 , 2022.

Jingtao Zhan, Jiaxin Mao, Yiqun Liu, Jiafeng Guo, Min Zhang, and Shaoping Ma. Learning discrete representations via constrained clustering for effective and efficient dense retrieval. In WSDM 2022 , 2022.

Tiancheng Zhao, Kyusong Lee, and Maxine Eskénazi. Unsupervised discrete sentence representation learning for interpretable neural dialog generation. In ACL 2018 , 2018.

Yujia Zhou, Jing Yao, Zhicheng Dou, Ledell Yu Wu, Peitian Zhang, and Ji rong Wen. Ultron: An ultimate retriever on corpus with a model-based indexer. ArXiv , abs/2208.09257, 2022.

Shengyao Zhuang, Houxing Ren, Linjun Shou, Jian Pei, Ming Gong, Guido Zuccon, and Daxin Jiang. Bridging the gap between indexing and retrieval for differentiable search index with query generation. ArXiv , abs/2206.10128, 2022.

Łukasz Kaiser, Aurko Roy, Ashish Vaswani, Niki Parmar, Samy Bengio, Jakob Uszkoreit, and Noam M. Shazeer. Fast decoding in sequence models using discrete latent variables. In ICML 2018 , 2018.

A Datasets details

We conduct experiments on three document 5 retrieval datasets: NQ320K, MS MARCO, and BEIR.

NQ320K. NQ320K is a popular dataset for evaluating generative retrieval models [37, 41]. It is based on the Natural Questions (NQ) dataset proposed by Google [15]. NQ320K consists of 320k query-document pairs, where the documents are gathered from Wikipedia pages, and the queries are natural language questions. We follow the evaluation setup in NCI [41] and further split the test set into two subsets: seen test , in which the annotated target documents of the queries are included in the training set; and unseen test , in which no labeled document is included in the training set.

Note that the NQ320K dataset we utilized has been pre-processed based on the NCI [41] and includes approximately 100k documents. This differs from the DSI approach [37], which processed a version of the NQ320K dataset containing about 200k documents. The distinction lies in the method used to remove duplicate documents. In our case, we eliminated duplicates by comparing document titles, whereas DSI employed URLs. For example, pages https://en.wikipedia.org//w/index. php?title=Statue_of_Liberty&amp;oldid=804877528 and https://en.wikipedia.org/ /w/index.php?title=Statue_of_Liberty&amp;oldid=834310497 are two versions of entity 'Statue of Liberty' . DSI NQ320K treats them as separate documents, whereas our implementation considers them as a single document. We have found that the content of different versions of the same entity's pages is usually almost identical, with only minor variations, often occurring in later parts of the document. Consequently, we believe that distinguishing between different versions of an entity is beyond the model's capabilities, and that using different versions of the same entity as negative examples in training may hurt model performance.

MS MARCO. MS MARCO is a collection of queries and web pages from Bing search. To create the document collections, akin to NQ320k and following [46], we sample a subset of original documents by retaining the top-1 document for each query. We evaluate the models on the queries of the MS MARCO dev set and retrieval on the sampled document subset. We did not split the dev set into seen and unseen because 84% of the queries in the MS MARCO dev set are unseen.

BEIR. BEIR is a collection of datasets for heterogeneous retrieval tasks. In this paper, we evaluate the models on 6 BEIR datasets, which include distinct retrieval tasks and document collections from NQ and MS MARCO: (i) BEIR-Arg retrieves a counterargument to an argument; (ii) BEIRCovid retrieves scientific articles about the COVID-19 pandemic; (iii) BEIR-NFC retrieves medical documents from PubMed; (iv) BEIR-SciFact retrieves scientific papers for fact-checking; (v) BEIRSciDocs retrieves citations for scientific papers; (vi) BEIR-FiQA retrieves financial documents. All the queries in BEIR test set are unseen [38].

We summarize the statistics of above datasets in Table 4.

B Baselines

The sparse retrieval baselines are as follows:

BM25 , uses the tf-idf feature to measure term weights; we use the implementation from http://pyserini.io/ .

DocT5Query , expands a document with possible queries predicted by a finetuned T5 with this document as the input.

The dense retrieval baselines are as follows:

DPR [13], a dual-encoder model using the representation of the [CLS] token of BERT.

ANCE [42], an asynchronously updated ANN indexer is utilized to mine hard negatives for training a RoBERTa-based dual-encoder model.

Sentence-T5 [22], a dual-encoder model that uses T5 to produce continuous sentence embeddings. We reproduce Sentence-T5 (ST5 for short) on our datasets, the model is based on T5-Base EncDec model and is trained with in-batch negatives.

GTR [23], a state-of-the-art dense retrieval model that pre-trains sentence-T5 on billions of paired data using contrastive learning.

Contriever [11], a dual-encoder model pre-trained using unsupervised contrastive learning with independent cropping and inverse cloze task.

And the generative retrieval baselines are as follows:

GENRE [5], an autoregressive retrieval model that generates the document's title. The original GENRE is trained on the KILT dataset [25] using BART, and we reproduce GENRE on our datasets using T5 for a fair comparison. For datasets without title, we use the first 32 tokens of the document as pseudo-title.

DSI [37], which represents documents using hierarchical K-means clustering results, and indexes documents using the first 32 tokens as pseudo-queries. As the original code is not open source, we reproduce DSI using T5-base and the docids of NCI [41].

SEAL [1] uses arbitrary n-grams in documents as docids, and retrieves documents under the constraint of a pre-built FM-indexer. We refer to the results reported by Wang et al. [41].

CGR-Contra [17], a title generation model with a contextualized vocabulary embedding and a contrastive learning loss.

DSI-QG [47], uses a query generation model to augment the document collection. We reproduce the DSI-QG results using T5 and our dataset.

NCI [41], uses a prefix-aware weight-adaptive decoder and various query generation strategies, including DocAsQuery and DocT5Query. In particular, NCI augments training data by generating 15 queries for each document.

Ultron [46], uses a three-stage training pipeline and represents the document as three types of identifiers, including URL, PQ, and Atomic.

C Performance on retrieving new documents

In this experiment, we investigate the impact of various document tokenization techniques on the ability of generative retrieval models to retrieve new documents. The generative models with different tokenization methods are trained on NQ320K data, excluding unseen documents, and are evaluated on NQ320K Unseen test set and BEIR-{Arg, NFC, SciDocs} datasets. For the baseline methods, which use rule-based document tokenization methods, the docids are generated for the target document collection using their respective tokenization techniques. In contrast, our proposed method uses a tokenization model to tokenize the documents in the target collection, producing the docids. However, our method may result in duplicate docids. In such cases, all corresponding documents are retrieved and shuffled in an arbitrary order. The results of this evaluation are summarized in Table 5.

Document tokenization methods that do not consider the semantic information of the documents, such as Naive String and Atomic, are ineffective in retrieving new documents without model updating. Methods that consider the semantic information of the documents, such as those based on title or BERTclustering, show some improvement. Our proposed document tokenization method significantly improves over these existing rule-based document tokenization methods. For instance, when the model trained on NQ - a factoid QA data based on Wikipedia documents - is applied to a distinct retrieval task on a different document collection, BEIR-SciDocs, a citation retrieval task on a collection of scientific articles, our proposed document tokenization model still showed promising results with an nDCG@10 of 12 . 3 , which is comparable to those models trained on the target document collection. This suggests that our proposed method effectively encodes the semantic information of documents in the docid and leads to a better fit between the docid and the generative retrieval model.

D Efficiency analysis

In Table 6, we compare GENRET with baseline models on MS MARCO (323,569 documents) in terms of memory footprint, offline indexing time (not including the time for neural network training), and online retrieval latency for different Top-K values. We have four observations: (i) The memory footprint of generative retrieval models (GENRE, DSI-QG, and the proposed model) is smaller than of dense and sparse retrieval methods. The memory footprint of generative retrieval models is only dependent on the model parameters, whereas dense and sparse retrieval methods require additional storage space for document embeddings, which increases linearly with the size of the document collection. (ii) DSI and GENRET take a longer time for offline indexing, as DSI involves encoding and clustering documents using BERT, while GENRET requires tokenizing documents using a tokenization model. Dense retrieval's offline time consumption comes from document encoding; GENRE uses titles hence no offline computation. (iii) The online retrieval latency of the generative retrieval model is associated with the beam size (i.e., Top-K) and the length of the docid. GENRET utilizes diverse clustering to generate a shorter docid, resulting in improved online retrieval speed compared to DSI and GENRE.

E Embedding visualization

20

20

40

40

Embedding at position 1

Embedding at position 2 Document embedding Codebook embedding

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAs one of the mainstream paradigms for document retrieval, generative approaches have enjoyed a steady growth of interest thanks to the recent thriving of large language models. Document tokenization is a crucial step in generative retrieval, which is rule-based in most existing methods, usually generalizing poorly to unlabeled documents. To this end, the paper presents a novel document tokenization learning framework GENRET that learns to tokenize a document into semantic docids in a discrete auto-encoding scheme. GENRET consists of a shared sequence-to-sequence-based document tokenization model, a generative retrieval model, and a reconstruction model, optimized in an end-to-end fashion using a well-designed progressive training scheme to stabilize the training process. Experimental results show that GENRET significantly improves over the baseline methods, especially in generalizing to unseen documents.\n\nSTRENGTHS:\n1. The authors have identified a significant problem in existing generative retrieval approaches. Conventional practice typically treats document tokenization as a fixed pre-processing step; however, this is not optimal as ad-hoc document tokenization methods often fail to capture the complete semantics of a document and generalize poorly to unlabeled documents. Therefore, the paper proposes for the first time to take document tokenization as a learnable module and introduces a novel framework to represent documents as discrete docids that effectively capture the semantic information of the document.   \n2. The auto-encoding scheme is not new but reasonably integrated into the proposed framework. The progressive training scheme effectively addresses the challenge of learning docids in an autoregressive fashion and stabilizes the training process. The diverse clustering techniques successfully facilitate the diversity of generated docids.   \n3. Extensive experiments on various benchmark datasets have demonstrated the superiority of the proposed method against established document retrieval approaches.\n\nWEAKNESSES:\n1. I think the presentation of the method description (Sec 3) could be further improved, given that some details are not clearly explained. For instance, according to Eq. 4, the reconstruction model is non-parametric, and the stop gradient operator is applied to $\\mathbf{d}_t$ and $\\mathbf{d}^*_t$, so which part of the parameters the reconstruction loss is aimed at optimizing? I guess it's probably the parameters of the codebook and the encoder-decoder. However, the authors mentioned that \"the gradients to document representation $\\mathbf{d}_T$ are defined as $\\frac{ \\partial \\mathcal{L}_r}{ \\partial \\mathbf{d}_T} \\coloneqq \\frac{ \\partial \\mathcal{L}_r}{ \\partial \\mathbf{z}_T}$ \" (in Line 173), but $\\mathbf{z}_T$ is actually taken from the codebook based on Eq. 5, so the reconstruction loss $\\mathcal{L}\\_{rec}$ is to optimize the codebook only? Moreover, does the generative retrieval model $P$ share the same codebook with the document tokenization model $Q$? What about the encoder and decoder? In addition, this section also does not cover the inference phase after the training completion. Finally, I do not understand why the docid re-assignment technique can promote the diversity of generated docids.   \n2. All experimental evidence shows that the proposed approach is effective on purely quantitative measures. However, I do not see any qualitative experiments presented, e.g., by visualization analysis to demonstrate that the codebook embeddings are uniformly scattered in the document representation space or the diversity of generated docids is better than baselines, etc.\n\nQUESTIONS:\n**Q1:** Please take a look at **Weakness1** for suggestions to improve the presentation.   \n**Q2:** In Eq. 8, is the final loss directly the sum of reconstruction loss, commitment loss, and retrieval loss? Should we set different weighting coefficients for each loss to balance the effects between them?   \n**Q3:** For the implementation details (Sec 4.3), the authors say that \"the length of the docid $M$ is dependent on the number of documents present. For datasets containing a larger number of candidate documents, a larger value of $M$ is set to ensure that all documents are assigned unique document ids\". So what $M$ is set to for each of the three datasets used in this paper? Is it possible to suggest how $M$ should be taken for a given number of documents from the perspective of theoretical analysis?\n\nLIMITATIONS:\nAs mentioned in the conclusion, the limitations of this work are two-fold. One is the lack of validation of the proposed method on large-scale data, and the other arises from the insufficient generalization of the model to unoptimized document collections in different types or domains.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nCurrent document retrieval systems map documents and queries to doc ids. These docids are assigned randomly, by clustering by using text/attribute information. The authors propose a learned document tokenization scheme where the the semantics of the documents are encoded into learned docids. These docids are used in end-to-end document retrieval methods. The proposed method consists of a retrieval model, a tokenization model and a reconstruction model. All three components are jointly trained. The authors show that on slighgt improvements on many different datasets. They in particular show that thir method generalized much better to unseen data.\n\nSTRENGTHS:\nRecently introduced end-to-end retrieval models are becoming more popular because they result in better performance. Learning better tokenization schemes is indeed a very relevant and important problem. The authors investigate a reasonable approach of using an auto-encoder to obtain doc ids. The 14% increase in performance on unseen data for the NQ dataset does suggest that the proposed method leads to learning better docids.\n\nWEAKNESSES:\nThe main weaknesses are as follows:\n* It is not clear why the autoencoder and the retrieval model are jointly trained. Rajput et. al. [1] also propose using an auto-encoder to learn docids. They use a RQ-VAE to learn docids that are conducive for auto-regressive generation, where each additional token in the docid can be thought of as the residual. In their method the auto-encoder and the retrieval model can be separately learned. \n* The evaluation using an unseen dataset is only done with the NQ dataset.\n* Most retrieval methods usually use a variant of contrastive loss or a cross-entropy loss. It is not clear why the authors use both and there is ablation for the retrieval loss.\n\n[1] Shashank Rajput, Nikhil Mehta, Anima Singh, Raghunandan H Keshavan, Trung Vu, Lukasz\nHeldt, Lichan Hong, Yi Tay, Vinh Q Tran, Jonah Samost, et al. Recommender systems with\ngenerative retrieval. arXiv preprint arXiv:2305.05065, 2023.\n\nQUESTIONS:\nHere are some questions I have: \n* How is the 'GENRET w/o learning' model trained? The paper mentions that the proposed learning scheme is not used for  'GENRET w/o learning.' Does this mean the progressive training scheme is not used? I'm trying to understand why GENRET performs so much better on the unseen test set. \n* Can you compare GENRET with [1]?\n* Have you inspected the learned docids? Do you see similar documents being assigned similar docids? Do you see meaningful clusters arise when you group by docids?\n* How important is the commitment loss?\n* How much better is GENRET when compared to other baselines on an unseen test set from other datasets? Is the improvement as significant as it is with the NQ dataset? The NQ test set is quite small and dividing it into seen and unseen further reduces the number of samples and the generalization might be more convincing if shown on larger datasets or on more datasets.\n\n[1] Shashank Rajput, Nikhil Mehta, Anima Singh, Raghunandan H Keshavan, Trung Vu, Lukasz\nHeldt, Lichan Hong, Yi Tay, Vinh Q Tran, Jonah Samost, et al. Recommender systems with\ngenerative retrieval. arXiv preprint arXiv:2305.05065, 2023.\n\nLIMITATIONS:\nYes, the authors to discuss limitations in terms of scale of dataset and generalization to different types/domains.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
GENRET, the proposed model learns to tokenize documents into short discrete representations (i.e., docids) via a discrete auto-encoding approach. Authors develop a progressive training scheme to capture the autoregressive nature of docids and diverse clustering techniques to stabilize the training process.

STRENGTHS:
- Very strong results on NQ, MS MACRO
- Beats dense retrieval.
- Nice ablation experiments.

WEAKNESSES:
-	Very handwavy in lines 173-174. Explanation on how argmax is bypassed is unintuitive
-	Doc id reassignment looks to depend upon the documents in the batch. Was any specific batching done?

QUESTIONS:
- GenRet w/o learning part in lines 297 unclear.
- The paper is very dense in math. Authors may consider writing out some more textual explanations of what is happening.

LIMITATIONS:
Limitations were properly addressed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

comments: SUMMARY:
This paper introduces GenRet, an auto-regressive retriever that focuses on finding the right clusters (or document ID or document tokenization) approach. Compared to previous generative retriever approaches, GenRet uses three different losses, progressive training, and clustering techniques. Overall, the paper is quite interesting and has the potential to improve the performance of information retrieval systems.

STRENGTHS:
Interesting and important problem to work on. The document representation/tokenization for generative retrievers is one of the key problems here. I am glad that there is solid research here.

WEAKNESSES:
Prior generative retrievers use a standard encoder-decoder with a pre-computed document ID. However, GenRet requires several modifications to work properly. While this is not necessarily a problem, I believe that some explanation and ablation studies are needed. For example, the retrieval loss makes GenRet a generalization of dual encoders. If the code-book size is the same as the number of documents in the candidate pool, and the length of the doc ID is 1, then GenRet will reduce to something very close to dual encoders. Therefore, I would like to see more analysis on how to decompose the doc ID and how much it affects generation. This type of study cannot be done in prior approaches because the doc IDs are fixed.

QUESTIONS:
Could you provide more insights in what way the generative doc id is better than prior pre-computed doc id? Did you find the doc id semantically meaningful?

LIMITATIONS:
The BEIR results are still behind non-generative retrievers.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper works on an emerging research direction, generative retrieval,  where retrieval is considered as a generating the document ids. This paper proposed to learn a seq2seq model to generate docids from the document. The challenge lies in how to propagate the retrieval loss (from another seq2seq model: query -> relevant docids) to the docid generation model.  The paper proposes several techniques, including commitment loss, iterative optimization, and document reconstruction. \n\nExperiments were done on NQ, MS MARCO and BEIR. Results show superior performance of the proposed method compared to prior generative retrieval approaches as well as standard dense retriever baselines.\n\nSTRENGTHS:\n- It is not straight forward to learn a latent docid from query-document relevance labels. The authors propose several techniques to make enable this. \n- Authors adopt diverse baselines \n- The paper is clearly written\n\nWEAKNESSES:\nMy main concern is around the MS MARCO results, which seems very different from other published results. E.g., in the ANCE paper, ANCE's MRR is 0.33 [1]; docT5query's MRR is 27.2 [2]. The MRR numbers reported in this paper is ~50-60, and docT5query is better than ANCE according to this paper which is not true based on prior research. This makes me question the soundness of this set of experiments.\n\nFor BEIR experiments, it is unclear how the six datasets are selected from all BEIR tasks.\n\n\n[1] Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval. Xiong et al. \n[2] https://cs.uwaterloo.ca/~jimmylin/publications/Nogueira_Lin_2019_docTTTTTquery-v2.pdf\n\nQUESTIONS:\n- Please explain why the MS MARCO results are off from results reported by prior work, and why dense retrieval (ANCE) is worse than docT5query on this set up.\n- How was the six BEIR dataset picked?\n\nLIMITATIONS:\nWould be nice to discuss how this method scales up to larger corpus, and whether / how it supports index updating.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-09"}