Learning to Tokenize for Generative Retrieval
Authors: Weiwei Sun
Lingyong Yan
Zheng Chen
Shuaiqiang Wang
Haichao Zhu
Pengjie Ren
Zhumin Chen
Dawei Yin
Maarten de Rijke
Zhaochun Ren
Venue: NeurIPS 2023 poster (2023)
Keywords: Information Retrieval
Document Retrieval
Generative Retrieval
TL;DR: This paper proposes a document tokenization learning method based on discrete auto-encoding to improve 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
Weiwei Sun', Lingyong Yan', Zheng Chen', Shuaiqiang Wang', Haichao Zhu? Pengjie Ren', Zhumin Chen', Dawei Yin?, Maarten de Rijke", Zhaochun Ren** +*
'Shandong University, China "Baidu Inc., China
University of Amsterdam, The Netherlands *Leiden University, The Netherlands yindavei@acm org
{sunnweivei, wwt.eiiieeiiinnnnnnnne.e Sagyyongy@gmmml.... m m.derijje@uva,
z. z.renllaccs Weidenuniv..n.
Abstract
As a new paradigm in information retrieval, generative retrieval directly generates ranked list of document identifiers (docids) for given query using generative language models (LMs). How 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- muse- tokenization, which is ad-hoc and does not generalize well. In contrast, in this paper we propose novel document tokenization learning method, GENRET, which learns to encode the complete document semanties into docids. GENRET learns tokenize documents into short discrete representations (i.e., docids) via a discrete auto-encoding approach. We develop progressive training scheme capture the autoregressive nature of docids and diverse clustering techniques to stabilize the training process. Based on the cemantic-embedded docids of any 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-ee 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.
Introduction
Document retrieval plays an essential role in web search applications and various downstream knowledge-intensvy 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 ranked list of document identifiers (docids) for 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]. also offers a promising approach to better exploit the capabilities of recent large LMs [1, 41].
As shown in Figure document tokenization aims to tokenize each document in the corpus sequence of discrete characters, i.e., docids. Document tokenization plays crucial role in generative retrieval, as it defines how the document is distributed in the semantic space [37]. And it still an open problem how to define docids. Most previous generative methods tend to employ document tokenizers, such as generating titles or URLs [5, 46), or clustering results from off-the thethf document embeddings [37, 41]. Such baled 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].
based
Corresponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Query
Transformer
Codebook
32
45
21
Generative Retrieval
docid z
Transformer
Codebook
32
45
21
R(dl2)
docid
Document Document Tokenization
Reconstruction
Reconstructed
Figure 1: An overview of our proposed method. The proposed method utilizes a document tokeniza- tion model to convert given document into sequence of discrete tokens, referred to as docid. This tokenization process allows for the reconstruction of the original document through reconstruction model. Subsequently, an autoregressive generation model is employed to retrieve documents through the generation of their respective docids.
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. GEN- RET consists of shared sequence-to-sequeee -es---eueeebbsseeeeeee document tokenization model, generative retrieval model, and a reconstruction model. In the proposed auto-encoding learning scheme, the tokenization model learns to convert documents 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 autoregressive manner for given query. The above three models optimized in an end-to-end fashion to achieve seamless integration.
There are usually two challenges when using auto-encoding to optimize 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, scheme allows for 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) reconstruction loss for predicting the
devise a progressive training scheme. This training
document using the generated docid, (ii) a commitment loss for committing the docid and avoiding forgetting, and (iii) retrieval loss for optimizing the retrieval performance end-to-end To address the second challenge, we propose parameter initialization strategy and a re-assignment of the docid based on 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 perfor- mance against state-of-the- generative retrieval models on NQ320K. GENRET achieves + 14% relative improvements 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.
the unseen test set of NQ320K over the best generative retrieval baseline.
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 tokenization model are reconstructed by 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, generative retrieval baselines.
2 Preliminaries
The document retrieval task can be formalized as the process of retrieving a relevant document d for search query q from a collection of documents D. Each document d E D is assumed to be a plain text where |d] is the total number of
consisting of a sequence of tokens, denoted as d = {di,.. ....... d tokens the document. For generative retrieval models, it is usually challenging and computationally
"The code of this work is available www. www.github..o comm/sumnveevee/Gentee.

---

inefficient to directly generate original documents of typically long length. Therefore, most existing approaches rely on technique named document tokenization, which represents document d= [d1, d|a) } as shorter sequence of discrete tokens (docid) each token 2, is as K-way categorical variable, with E [1, 2,...... ...,...,, ,K], and M is the length of the docid. See Figure I for example of document tokenization with A/ and K 64.
[=1,...: ....... =M ), where
As an alternative sequence of the original document, the tokenized docid & 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 is sequence of 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 - to map to docid More details about are provided in Section 3.1. After tokenizing each document docid 2, generative retrieval model P: + learns to retrieve relevant documents by generating a query to a docid 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, 37J. However, 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 auto- encoding, learn semantic docid in fully end-to-end manner. Figure 1 gives overview of the proposed method. The proposed GEN RET comprises three main components: (i) sequence-to- sequence based retrieval model P(z q). (ii) 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 2, and the retrieval model is trained generate the latent variables 2 for given query q. In addition, the reconstruction model is used to re- e--enerrtee the original document from 2 to ensure 2 captures the semantics of the original document much as possible.
3.1 Model architecture
Following DSI [37], we employ an encoder-decoder Transformer to implement the generative retrieval model. Specifically, given input text d² the T5-based tokenization model encodes d and prefix of docid ze, and continuously produces latent representation d, of d at time step t:
d, = Decoder (Facoder(a z2() E R.P,
where D denotes the hidden size of the model. Encoder(d) (d) denotes the output of the Encoder.
Then, the tokenization model generates a token for each document based on d,. At each timestep f, we define an external embedding matrix named codebook E, E IR= XxD, where K is the size of the discrete latent space. There are K embedding vectors er,j RR ..j E [K], and each vector e can be regarded as the centroid of a segmentation.
Based on the codebook E,, the discrete latent variable 2 at timestep t is calculated by a dot-producc look-up using the codebook E,:
Q(2 =j ser, d) Softmax, E.).
where Q(2 = 2<0, d) denotes the probability of tokenizing d to particular value j E [K] at timestep t, Softmax, is 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, propose auto-encoding training scheme, where a reconstruction model R: 2 + that predicts d using & is designed to force the tokenization model Q: d + reproduce docid that can be reconstructed tack-to-the original document.
We use document d here for the denotation, noting that the computation is the same when y is input.
3

---

The input of the reconstruction model is docid and the output is its associated document d. We first embed 2 into representation matrix z (ZI....,ZM. E R RMDD using the codebook of the
tokenization model:
(e1,x) 2,221 eM R
(3)
where each E [M] = 0t,., E R is the embedding vector of z, in the 1-step codebook E,,
We then devise retrieval- based reconstruction model that predicts the target document d by retrieving from document collection D, based the inputs z. The relevance score between the input docid and the target document d is defined as follows:
exp(zt sg(d!)) ==1 Edesteed exp(zi sg(d+,))
M II
R(d z)
(4)
where S(Ec)) is sub-collection of 'D consisting of documents that have docid prefix that is the same as Zez. d² E S(z<z) represents a document from the sub-collection S(zzt). d, and d'; are continuous representations of documents d and d" respectively, defined in Eq. 1. The operator sg(+) is the stop gradient operator to prevent gradient back propagation. Intuitively, Rid z) is designed to retrieve specific document from a set of documents S(221) at each timestep t. The set ((2ce) 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 Ex(--
3.2 Model optimization
For the document tokenization model Q(z d), generative retrieval model P(z g). and reconstruc- tion model R(d =). jointly optimizing these three models using auto-encoding is challenging due to the following reasons: (i) Learning docids an autoregressive fashion. On one hand, prediction of the at time relies previously predicted docids Rer, which is often under-opti- mized at the beginning and rapidly changes during training, making it difficult to reach convergence. On the other hand, simultaneously optimizing makes it challenging 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 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 diverse clustering techniques ensure docid diversity (see Section 3.2.2).
guarantee unique docid
3.2.1 Progressive training scheme
To optimize each of the three models listed above in an autoregressive manner, propose progressive auto- encoding 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 2T step T' E [M]] is learned and optimized at the corre- sponding learning step. Besides, at each step T' E [M]. the docid 2T and the model parameters associated with 2T generation are updated, while previously produced docids Zcr and other parameters are kept fixed. By progressively performing the above process, we can finally optimize and learn our models.
Document
GenRet
Fixed docid
Update docid Unlearned dockd
32 45
21
Progressive training scheme. 2e (docid timestep 1) is optimized the I-th training step, while zez. (docids before timestep 1) are kept fixed.
Figure 2:
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 =) as an auxiliary model to learn to optimize the docid generation, whose main goal is capturing much semantics in the docid as

---

possible. Therefore, we define a reconstruction loss function of step T as follows:
CRee = log R(d| ågT); where ägT (sg(,)).....,58 E [T] 24 Cr E RD, where j = arg max ( = kcr,
E RT
(5)
where žcr is the first T representations of the 2, and only the variable žr is optimized in step T. Q(= = zee,d) d) is defined in Eq. 2. And the document tokenization model Q can therefore be optimized when minimizing PRee.
Of note, since the computation involves a non-differentiable operation (arg max(:)). we apply straight-hrouuhh gradient estimation paok-prop the gradient from reconstruction loss to dr following [39, 44], which copies the gradient of zr directly to dr. Specifically, the gradients to document representation dr are defined as dLaz embedding er,j are defined as
And the gradients to the codebook Dan
ad, &C
GET
8C*. zy= bsr Ber.j
Commitment loss. In addition, to make sure the predicted docid commits to an embedding and to avoid models forgetting previous docid Bee, we add commitment loss as follows:
Lcom =
Σ log Q(=+ 2et, d).
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 query g. Specifically, suppose (g, d) are query and relevant document pair; we define the learning objective of retrieval model P as:
exp(dr dr)
Lrat - -log
log P(zx 2 4)2
Ed exp(dr d'r)
where the first term is a ranking-orienteen loss enhancing the model using (q, d) pair; d' is an in-batch document sampled from the same training mini-batch B; qr and dr denote the representation of g 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 step-T is the sum of reconstruction loss, commitment loss, and retrieval loss:
= Cree + LCom LRa.
(8)
3.2.2 Diverse clustering techniques
To ensure diversity of generated docids, we adopt two diverse clustering techniques-codebol) 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 docid assignments.
Codebook initialization. In order to initialize the codebook for our model, we first warm-up the model by passing the continuous representation dy the reconstruction model instead of the docid representation Zy as defined in Eq. 3. During this warm-up phase, we optimize the model using the reconstruction loss LRee and commitment loss LCam- Next, we collect the continuous representations dr of all documents in DD and cluster them into K groups. The centroids of these clusters are then used as the initialized codebook E,. To balance the initialized docid distribution, we utilize a diverse constrained clustering algorithm, Constrained K- 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 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, = {a.") ..... d (B) IR Bx denote the continuous representation of a batch of documents with batch size of B. The dot-product results are represented by H D. E, E RBxK To obtain distinct docids, calculate an alternative H Diag(u) (1) exp( H) Diag(v), where and
re-normalization vectors in IR and R respectively. The re-normalization vectors 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 24.
5

---

4 Experimental Setup
4.1 Datasets and evalutaion metries
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 lest, based on whether the target documents of the query have annotated queries in the training data, to test the gencralization 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 @ 1110,100) } and Mean Reciprocal Rank (MRR) @ 100 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
We consider 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-Contre [17], DSI- QG [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 for more details on the other baselines.
4.3 Implementation details
Hyper-paree prreepparameters. In our experiments, we utilize the T5-Base model [27] as the base Transformer and initialize new codebook embedding E, for each time step. The parameters of both the encoder- decoder 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 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 parpmm-prreeter e is set to 1.0, and the Sinkhorn- -nnhop 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 DocTSQuery [24] to augment the NQ and MS MARCO datasets. In query generation, we use nucleus sampling with parameters five queries for each document in the collection. For the BEIR datasets, we use the queries generated by GPL [40]. GPL uses 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].
0.8, f 0.8 and generate
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 Adam W and set the learning rate to 5e - 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 factor of 0. 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* obtain diverse clustering results. During inference, we use beam search with constrained decoding [5] and beam size of 100.
constrained

---

Table I: Results on Natural Questions (NQ320K). The results of the methods marked with are from our own re-implementation, others are from their official implementation. and ** indicate significant improvements over previous-best generative retrieval baselines with p-value < 0.05 and p-value 0.01, respectively. and indicate significant improvements over previous-best dense retrieval baselines with p-value < 0.05 and p-value < 0.01, respectively. The best results for each
metric indicated in boldface.
Full test
Seen test
Unseen test
Method
R @
R@ 10
R@ 100
MRR R@I
R@10
R@ 100 MRR R@1
R@I0
10
R@ 100
MRR
Sparse retrieval
BM25 [32]
29.7
60.3
82.1
40.2
29.1
59.8
82.4
39.5
32.3
61.9
81.2
42.7
DocTSQuery [24]
38.0
69.3
86.1
48.9
35.1
86.4
72.9
85.0
57.0
Dense retrieval
DPR [13]
50.2
77.7
90.9
59.9
50.2
78.7
91.6
60.2
50.0
74.2
88.7
58.8
ANCE [42]
50.2
78.5
91.4
60.2
49.7
79.2
92.3
60.1
52.0
75.9
88.0
60.5
Sentence-TSS [22]
53.6
83.0
93.8
64.1
53.4
83.9
94.7
63.8
56.5
79.5
90.7
64.9
CTR-Base [23]
56.0
84.4
93.7
66.2
54.4
84.7
94.2
65.3
61.9
83.2
92.1
69.6
Generative retrieval
GENRE [5]
55.2
67.3
75,4
59.9
69.5
83.7
90.4
75.0
6.0
10.4
23.4
7.8
DSI [37]
55.2
67.4
78.0
59.6
69.7
90.5
74.7
1.3
31.5
SEAL [1]
59.9
81.2
90.9
67.7
CGR-Contra [17]
63.4
81.1
DSI-QG" [47]
80.7
88.0
69.5
68.0
85.0
91.4
74.3
45.9
65.8
76.3
52.8
NCI [41]
66.4
85.7
92.4
73.6
69.8
88.5
94.6
76.8
54.5
75.9
84.8
62.4
Ours
68.1
88.8
95.2"
75.9
70.2
90.3*
96.0
77.7"
62.5
83.6"
92.5"
70.4
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 new state-of-the-art reveal that existing generative retricval methods perform well on the seen test but lag behind dense retrieval methods the unseen test. For example, NCI obtains MRR @ 100 of 76.8 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 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.
the NQ320K dataset. Furthermore, our results
Results on MS MARCO. Table presents the results on the MS MARCO dataset. GENRET outperforms generative retrieval methods such Ultron and dense retrieval baselines such ANCE and Se. Furthermore, previous generative retrieval methods (e.g., GENRE, Ultron) utilizing metadata such as the title and URI., while exhibiting decent performance on the NQ320K dataset, underperform comparison 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 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 lists the results of the baselines and GENRET on six datasets 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 some datasets, such as BEIR-Covid and BEIR- SciDocs. This may be because the titles of the documents these datasets do not adequately capture their semantic content.

---

Table 2: Results on MS MARCO. The re- sults of the methods marked with are from our own re-implementation. other results are cited from the original paper or imple- mented using official code. The best results are indicated in boldface.
Table 3: nDCG@10 results on BEIR. The re- sults 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-T's T3 trained on GPL datasets [40].
Method
@@1 R@ R@ 100 MRR
Method
Arg Covid NFC SciF. SeiD. FIQA Avg.
Sparse retrieval
Sparse retrieval
BM25 [32]
39.1
69.1
86.2
48.6
BM25 [32]
29.1 58.9
33.5
67.4
14.8
23.6
37.8 29.1 41.9
DocT5Query [24]
46.7
76.5
90.4
56.2
DocT5Query [24] 34.9 71.3
3288
67.5
16.2
Dense retrieval
Dense retrieval
ANCE [42]
45.6
75.7
89.6
55.6
ST5 GPL! [22]
32.1
74.4
30.1
58.6
12.7
26.0
39.0
Sentence-TS [22]
41.8
75.4
91.2
52.8
Contriever [11]
40.0
68.8
33.5
61.4
16.3
30.7
41.8
Generative retrieval
Generative retrieval
GENRE [S]
35.6
S7.6
79.1
42.3
GENRE [47]
42.5
14.7
20.0
42.3
6.8
11.6
30.0
Ultron-URL [46]
29.6
67.8
40.0
Ours
34.3
71.8
31.6
63.9
14.9
30.2
41.1
Ultron-PQ [46]
31.6
73.1
45.4
Ultron-Atomic [46]
32.8
74.1
46.9
Ours
47.9
79.8
91.6
58.1
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 box with docid and a diversity metric d, which is calculated by: d 1 1, K documents, n, denotes the number of documents that have a docid and nu, is the expected number of documents per docid under the uniform distribution.
Σ)=1 In, naul, where |-| represents the absolute value, n denotes the total number of
The results demonstrate the superiority of GENRET (represented by the yellow line) in terms of distribution uniformity. uses all the potential docid = 512 and achieves the highest diversity metric with value of d = 0.90. The method without docid reassignment also yields 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 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 demtnstrates 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 small model, and its performance is inferior to that of GENRET using T5-Base. However, the gap between the performance 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.
8

---

Frequency
Recall@1
5%
GanRat
75
GanRat Reassigmmnt
Sean
4%
GenRet Intialiation
Unseon
GenRet we Diverse Cluster
3%
512 d 0.90
2%
0.00
65
34
A 494
1%
2-0.77
if 0.77
0%
45
100
200
300
400
500
GenRet
GenRet
GenRet
Docid
learning
T5-Small
Figure 3: Left: Docid distribution on NQ320K. The id are sorted by the assigned frequency. Right:
Ablation study on NQ320K.
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 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 second-level docid 173 is added, the corresponding documents become more specifically related to email. With the addition of a third-level docid 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 hierarchical semantic structure within the learned docids.
Figure: 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 more challenging as the semantic granularity becomes finer.
Doc ID
Documant Trtle
stamp
email
Used mai.
omanl
173
64B
Mail merge
Docidt 388
Dockd 338 173
Docid 178-1
DGB
External mall preny
379 -992
Business letter
Conten't hide
608 222
US Poozal Service
Docid 204
502
Dockd 502
Docld:
te season
present
FINC
186 140
Mother Tracy McConnel
tba
Icatiliattoo
Figure 4: Left: Document titles along with their corresponding docids. It is observed that documents with similar docids tend have more relevant content. Right: Word cloud representing documents grouped by docid prefixes. This illustrates that different positions of the docid correspond to different levels of information, and the semantics within each cluster are closely related.
6
Related work
Sparse and dense retrieval. Traditional sparse retrieval calculates the document score using term matching metrics such as TE-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 pre-
trained 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 improve DSL. 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 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- unto-enc 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 's representation of the two tasks and are optimized in an end-to-end manner.
Discrete representation learning. Learning discrete representations using neural networks 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. Dall- E [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 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 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. 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 collec- tions in different types or domains, as compared 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 investigate the dynamic adaptation of docid prefixes for progressive training.
10

---

Acknowledgements
This work was supported by the Natural Science Foundation of China (62272274, 61972234, 62072279, 62102234, 62202271), the Natural Science Foundation of Shandong Province (ZR2021QF129, ZR2022QF004), the Key Scientific and Technological Innovation Program of Shandong Province (2019122Y010129, the Fundamental Research Funds of Shandong Univer- sity, the China Scholarship Council under grant nr. 202206220085, the Hybrid Intelligence Center, 10-year program funded by the Dutch Ministry of Education, Culture and Science through the Netherlands Organisation for Scientific Research, https: / //hybriddd-- mantelligence--ee centre....... .nl, and project LESSEN with project number NWA.1389. 1389.20.182 of the research program NWA ORC 2020/21,, which (partly) financed by the Dutch Research Council (NWO), All content represents the opinion of the authors, which is not necessarily shared endorsed by their respective employers and/or sponsors.
References
[1] 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.
[2] Paul S. Bradley, Kristin P. Bennett, and Ayhan Demiriz. Constrained k-means clustering. Microsoft Research, 2000.
[3] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan. Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Vosse Gretchen Krueger,
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 NeurlPS 2020, 2020.
[4] 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/161 111661100920 .09268, 2016.
[5] Nicola De Cao, Gautier Izacard, Sebastian Riedel, and Fabio Petroni. Autoregressive entity retrieval. In ICLR 2021, 2021.
[6] 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.
[7] Jiangui Chen, Ruqing Zhang, Jiafeng Guo, Yigun Liu, Yixing Fan, and Xueqi Cheng. Corpus- Brain: Pre-train a generative retrieval model for knowledge-intens intensive language tasks. In CIKM 2022, 2022.
[8] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In NIPS 2013, 2013.
[9] Daniel Gillick, Alessandro Presta, and Gaurav Singh Tomar. End-to-end retrieval continuous space. ArXiv, abs/181 0tss/8811.0008.. 2018.
[10] Sebastian Hofstätter, sheng-Chiel 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.
[11] 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.
[12] Shuning Jin, Sam Wiseman, Karl Stratos, and Karen Livescu. Discrete latent variable represen- tations for low-resource text classification. In ACL 2020, 2020.
11

---

[13] Vladimir Karpukhin, Barlas Oğuz, 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.
[14] Omar Khattab and Matei Zaharia. ColBERT: Efficient and effective passage search via contex- tualized late interaction over BERT. In SIGIR 2020, 2020.
[15] 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: benchmark for question answering research. TACL, 7:453 466, 2019.
[16] John Lafferty and ChengXiang Zhai. Document language models, query models, and risk minimization for information retrieval. In SIGIR 2001, 2001.
[17] 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.
[18] Jimmy Lin, Rodrigo Nogueira, and Andrew Yates. Pretrained Transformers for Text Ranking: BERT and Beyond. Springer Nature, 2022.
[19] 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.09 0991553,, 2022.
[20] Sanket Vaibhav Mehta, Jai Gupta, Yi Tay, Mostafa Dehghani, Vinh Quang Tran, Jinfeng Rao, Marc-Alexander Najork, Emma Strubeli, and Donald Metzler. DSI++: Updating transformer memory with new documents. ArXiv, abs/2212.097444, 2022.
[21] Donald Metzler, Yi Tay, and Dara Bahri. Rethinking search: Making domain experts out of dilettantes. In SIGIR 2021, 2021.
[22] Jianmo Ni, Gustavo Hernández Ábrego, Noah Constant, Ji Ma, Keith B. Hall, Daniel Matthew Cer, and Yinfei Yang. Sentence-TSy T5: Scalable sentence encoders from pre-trained text-to-text models. In Findings of ACI. 2022, 2022.
[23] 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.
[24] Rodrigo Nogueiral and Jimmy Lin. From doc2query to doctiiiquuer. attps://cs. uwaterloo.. ca/" pdf, 2019.
[25] 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.
[26] Yinggi 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 NAACI 2021, 2021.
[27] 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. JMIR, 2020.
[28] 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 NewrIPS 2023,
2023.
12

---

[29] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Iya Sutskever. Zero-shot text-to-image generation. In ICML 2020, 2020.
[30] Nils Reimers and Iryna Gurevych. Sentence- Sentence embeddings using siamese bert- networks. EMNLP 2019, 2019.
[31] Stephen E. Robertson and Steve Walker. On relevance weights with little relevance information. In SIGIR 1997, 1997.
[32] Stephen E. Robertson and Hugo Zaragoza. The probabilistic relevance framework: BM25 and beyond. Found. Trends Inf. Retr., 2009.
[33] Jason Tyler Rolfe. Discrete variational autoencoders. In ICLR 2017, 2017.
[34] 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.
[35] Weiwei Sun, Lingyong Yan, Xinyu Ma, Shuaiqiang Wang, Pengjic 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.
[36] Yubao Tang, Ruqing Zhang, Jiangui Guo, Jiangui Chen, Zuowei Zhu, Shuaiqiang Wang, Dawei Yin, and Xueqi Cheng. Semantic-enhance differentiable search index inspired by learning strategies. In SIGIR 2023, 2023.
[37] 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 differentiable search index. In NeurIPS 2022, 2022.
[38] 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 NeurlPS 2021, 2021.
[39] Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In NIPS 2017, 2017.
[40] Kexin Wang, Nandan Thakur, Nils Reimers, and Iryna Gurevych. GPL: Generative pseudo labeling for unsupervised domain adaptation of dense retrieval. In NAACL 2022, 2022.
[41] 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 NeurlPS 2022, 2022.
[42] 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.
[43] 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.
[44] 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.
[45] Tiancheng Zhao, Kyusong Lee, and Maxine Eskénazi. Unsupervised discrete sentence represen- tation learning for interpretable neural dialog generation. In ACL 2018, 2018
[46] Yujia Zhou, Jing Yao, Zhicheng Dou, Ledell Yu Wu, Peitian Zhang, and Ji rong Wen. Ultron: An ultimate retriever on corpus with model-based indexer. ArXiv, abs/2208.092 009507, 2022.
[47] 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.
[48] Fukasz Kaiser, Aurko Roy, Ashish Vaswani, Niki Parmar, Samy Bengio, Jakob Uszkoreit, and Noam M. Shazeer. Fast decoding sequence models using discrete latent variables. ICMI. 2018, 2018
13

---

A Datasets details
Table 4: Statistics of datasets used in our experiments. The three values split by / on # Test queries denote the number of queries in the full, seen subset, and unseen subset, respectively, In BEIR, all queries in the test set are unseen.
Dataset NQ320K 109,739 7,830
Docs
# Test queries
# Train pairs
6,075 / 1,755
307,373
MS MARCO 323,569 5,187
807 4,380
366,235
BEIR-Arg 8,674
1,406
BEIR-CoVir 171,332
50
BEIR-NFC 3,633
323
BEIR-SciFact 5,183
300
BEIR-SciDoce 25,657
1,000
BEIR-FiQA 57,638
648
We conduct experiments three document retrieval datasets: NQ320K, MS MARCO, and BEIR.
NQ320K. NQ320K is popular dataset for evaluating generative retrieval models [37, 41J. 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 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-procee 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/ http///w/ndd..//e/eeeeeeeeeeeees iinexx.d.xx///ie pup"titlo=Statue_.ff of Lttoofff_ibrrtrtaaepaee wld.od /w/index. www.:ppp--tttttttteeeeeee
and https://en wikipedia org/ 01111-3331111111191 are two versions of entity "Statue of Liberty". DSI NQ320K treats them as separate documents, whereas our implementation considers them as single document. We have found that the content of different versions of the same entity's '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 NQ320k and following [46], we sample 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 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 BEIR datasets, which include distinct retrieval tasks and document collections from NQ and MS MARCO: (i) BEIR-Arg retrieves counterargument an argument; (ii) BEIR- Covid retrieves scientific articles about the COVID-19 pandemic; (iii) BEIR-NFC retrieves medical documents from PubMed; (iv) SciFFte retrieves scientific papers for fact-checking; (v) BEIR- SeiDoos retrieves citations for scientific papers; (vi) BEIR-FiQA retrieves financial documents. All the queries in BEIR test set unseen [38].
We summarize the statistics of above datasets in Table 4.
'A document in our context typically corresponds a webpage, while a passage is piece of text within a document. Our study focuses on building docid for documents. Using document retrieval tasks allows us to make fair comparisons with existing methods that utilize document URLs, titles, etc. And it also aligns with the settings of NQ32UK and BEIR. Passage retrieval, a different task, could lead to different evaluation results.
14

---

B Baselines
The sparse retrieval baselines are as follows:
BM25, uses the tf-idf feature measure term weights; the implementation from attp://pyserini. io/.
DocT5Query, expands document with possible queries predicted by a finetuned T5 with this document the input.
The dense retrieval baselines are as follows:
DPR [13], a dual-encoder model using the representation of the [CLS] token BERT.
ANCE [42], an asynchronously updated ANN indexer is utilized to mine hard negatives for training RoBERTa-baseee dual-encoder model.
Sentence-T5 [22]. dual-encoder model that uses T5 to produce continuous sentence
embeddings. We reproduce Sentence-TSs (ST5 for short) on our datasets, the model is based T5-Base EncDec model and is trained with in-batch negatives.
GTR [23], 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 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 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 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 contrastive learning loss.
DSI-QG [47], uses query generation model to augment the document collection. We reproduce the DSI-QG results using T5 and our dataset.
NCI [41], uses prefix- aware gies, including DocAsQuery and DocTSQuery. In particular, NCI augments training data by generating 15 queries for each document.
decoder and various query generation strate-
Ultron [46], uses three-stage training pipeline and represents the document as three types of identifiers, including URL,, PQ, and Atomic.
C Performance retrieving new documents
In this experiment, we investigate the impact of various document tokenization techniques on the ability of generative retrieval models 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. SciDoes) datasets. For the baseline methods, which use based document tokenization methods, the docids are generated for the target document collection using their respective tokenization techniques. tokenization model to tokenize the documents in the target collection, producing the docids. However, our method may result duplicate docids. such cases, all corresponding documents are retrieved and shuffled in arbitrary order. The results of this evaluation summarized in Table 5.
contrast, our proposed method uses
15

---

Table 5: Zero-shot evaluation on retrieving new documents with different document tokeniza- tion methods. The second column indicates the type of docid, where BERT-HC denotes BERT- Mieearchical-Cllsteriin [37], Prefi--HC denotes Prefix-aware BEERT-HieearchiccllCCCttttii [41],
and dAE denotes discrete auto-encoding.
NQ (R@1)
BEIR (nDCG@ 10)
Method
Docid
Unseen
Arg
NFC
SciDoes
DSI-Naive [37]
Naive String
0.0
0.1
1.0
0.1
DSI-Atomic' [37]
Atomic
0.0
0.2
0.8
0.1
GENRE [5]
Title
6.0
0.0
2.4
0.6
DSI [37]
BERT-HC
1.3
1.8
11.1
5.9
NCI [41]
Prefix-HC
15.5
0.9
4.3
1.2
Ours
dAE
34.2
12.1
12.1
12.3
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 BERT clustering, show some improvement. Our proposed document tokenization method significantly improves over these existing rule-based 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 different document collection, BEIR-SciDocs a citation retrieval task 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 better fit between the docid and the generative retrieval model.
Table 6: Efficiency analysis.
Method
Memory
Time (Offline)
Top-K
Time (Online)
ANCE
1160MB
145min
100
0.69s
GTR-Bast
1430MB
140min
100
1.97s
100
1.41s
GENRE
851MB
0min
10
0.69s
100
0.32s
DSI
851MB
310min
10
0.21s
100
0.16s
Ours
860MB
220min
10
0.10s
D Efficiency analysis
In Table 6, 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) 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 longer time for offline indexing, as DSI involves encoding and clustering documents using BERT, while GENRET requires tokenizing documents using tokenization model. Dense retrieval's '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.
16

---

E
Embedding visualization
Embedding at position
Embedding at position
Embedding at position
Document embedding
Codebook embedding
Figure 5: t-SNE visualization of the codebook embedding and document embedding on the NQ320K
dataset. The codebook embedding uniformly scattered in the document representation space.
17


Reviews:
[{"review_id": "YkRl6Snfiv", "reviewer": "Reviewer_zz1Z", "summary": "This 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.", "questions": "- 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?", "limitations": "Would be nice to discuss how this method scales up to larger corpus, and whether / how it supports index updating.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- 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", "weaknesses": "My 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your time and insightful comment. We would like to address your questions in turn.\n\n**Two reasons for different numbers for ANCE/docT5query**\n\n- First, please note that we are focusing on the document retrieval task in our paper, whereas the results you referenced are from the passage retrieval task. \n\n- Second, our setting is consistent with existing generative retrieval studies [1,2], where we retain the top-1 document for each query within our document collection. In contrast, the papers you cited utilize the MS MARCO leaderboard passage collection, which maintains the top-10 Bing Search passages for each query [3].\n\n\n**Higher results of docT5query**\n\n- We accredit the differences in results to the different retrieval tasks, and the size and construction methods of doc/passage collection. To provide more implementation details, for ANCE we use the [ance-firstp](https://huggingface.co/sentence-transformers/msmarco-roberta-base-ance-firstp). For docT5query we use [doct5query](https://huggingface.co/castorini/doc2query-t5-base-msmarco).\n\n\n**BEIR data selection**\n\n- Our selection of the six datasets from BEIR is based on the corpora size. In particular, the document size of selected data is around or less than 320K documents, intending to test the method's effectiveness on moderate-sized corpora following previous studies [1,2].\n\n**Scaling corpus and updating index**\n\n- Thanks for the questions! As we have discussed in conclusion, scaling the proposed model to larger corpora typically requires longer docid length, more model capacity (parameters), and more computation resources. In addition, expanding to a larger corpus may confronts several known challenges [4], such as knowledge forgetting and training efficiency, which need to be addressed in our future research.\n\n- We anticipate that the proposed semantic docid generation method will indeed be conducive to large corpora. This is because a more meaningful docid results in a more efficient compression of the document's semantics, thereby reducing the burden on the model to memorize the sequence of these IDs. In comparison, when the docid is arbitrary, eg native string, the model needs to memorize these ID sequences without the aid of meaningful association, leading to additional overhead. We observe relevant evidence in the rate of training loss decline, e.g., the loss declines faster when using a meaningful docid.\n\n- As for updating the index, we presume that expanding docid could provide a feasible solution e.g., adding new docid bits for new documents to differentiate them from existing ones, and continue training the model based on progressive training. In addition, the MOE approach might also be feasible, e.g., indexing documents in separate experts and merging the results of multiple experts through aggregation techniques. We would like to discuss this further and do more exploration in our final version.\n\nWe hope our response can address your concerns.\n\n[1] Transformer Memory as a Differentiable Search Index.\n[2] Ultron: An Ultimate Retriever on Corpus with a Model-based Indexer\n[3] MS MARCO: A Human Generated MAchine Reading COmprehension Dataset.\n[4] How Does Generative Retrieval Scale to Millions of Passages?"}, {"review_id": "OGCNstSyOt", "reviewer": "Reviewer_e88A", "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.", "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?\n", "limitations": "The BEIR results are still behind non-generative retrievers. \n", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\n", "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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your time and insightful comment. We would like to address your questions in turn.\n\n**Compare to dual-encoder**\n\n- Thanks for your insightful comment. We further compared the models with different values of M and K. The results are in the following table. We find that the performance of models with different configurations are close to the NQ320K. When length of the doc ID is 1, the model will be similar to dual encoder, but with the difference that doc embedding is not obtained by encoding with a doc encoder, but is obtained by training as a codebook parameter.\n| Metric | Model 1 |  Model 2 | Model 3 | \n| --- | --- | --- | --- | \n| K | 109,739 | 2,048 | 512 | \n| M | 1  |  2  |  3  |\n| R@1 |  68.5 | 68.8 |  68.1  |\n\n- Employing a longer docid will reduces model parameters but necessitates more training steps. Utilizing a codebook of length 1 significantly increase the numbers of model parameters, rendering the model unsuitable for large-scale data sets. As advised, we will be incorporating more comprehensive analyses to our final paper.\n\n**Qualitative analysis**\n\n- Thanks for your suggestions. In our newly uploaded PDF, we add two qualitative analysis:\n\n- Figure 1 lists document content and the corresponding docid generated by GenRet on NQ320K dataset, and shows that those documents with more similar docid have more relevant content. For example, we find that documents with docid started with 338-173 are related to Email, while documents with docid started with 338 are related to information exchange methods. Therefore, we find the generated docid semantically meaningful. \n\n- Figure 2 visualizes the codebook embedding and document embedding. We see that the codebook embedding does uniformly scatter in the document representation space.\n\n- Figure 3 illustrates the word cloud of documents grouped by GenRet produced docid prefix. We can see that different positions of docid represent different levels of information, and the semantics within each cluster are related.\n"}, {"review_id": "w0FsNIHn80", "reviewer": "Reviewer_BPS1", "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.", "questions": "- GenRet w/o learning part in lines 297 unclear.\n- 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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Very strong results on NQ, MS MACRO\n- Beats dense retrieval.\n- Nice ablation experiments.", "weaknesses": "-\tVery handwavy in lines 173-174. Explanation on how argmax is bypassed is unintuitive\n-\tDoc id reassignment looks to depend upon the documents in the batch. Was any specific batching done?\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive comments! We appreciate your feedback and will address each of your points in turn.\n\n**Line 173-174** \n\n- Thanks for your comments and we will add more explanation and clarification in this part in our final paper. Specifically, we employ straight-through gradient estimation that allows for the transmission of the $z_T$ gradient to $d_T$ following *Oord et al. 2017 [37]*, which copies the gradient of $d_T$ directly to $z_T$. We will add more explanations to this section and revise the presentation as you suggested.\n\n**Batching** \n\n- Thanks for your insightful question! 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 will clarify further on this matter in the final paper for better clarity.\n\n**W/o learning**\n\n- This variant directly employs the final docid assigned by the proposed system and trains a separate DSI model following the practice of DSI-QG. This variant demonstrates that the generative retrieval model jointly trained with auto-encoding objectives can represent documents more sensibly based on semantics. This could potentially 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. We will add more explanations of this variant in our final paper.\n\n**Presentation suggestion**\n\n- Thanks for your valuable comment. We will add more explanations in the final paper to increase readability.\n\n[37] Oord, Aäron van den et al. Neural Discrete Representation Learning."}, {"review_id": "wKChMPQ5dn", "reviewer": "Reviewer_BkMu", "summary": "Current 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.", "questions": "Here 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.", "limitations": "Yes, the authors to discuss limitations in terms of scale of dataset and generalization to different types/domains.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Recently 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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful review! We will address each of your points in turn. \n\n**Compared with Rajput et al.**\n\n- Thanks for your insightful comment. First, we believe it is beneficial to joint the modeling of tokenization and retrieval tasks on text retrieval tasks. From the perspective of task characterization both tokenization and retrieval tasks aim to convert textual semantics into docid within a shared discrete space.\nThus sharing parameters can potentially align the model's representation of the two tasks and improve their respective capabilities. From an empirical perspective, the results of the w/o learning variant that only employs docid' results for training a separate retrieval model reveal the merits of joint modeling. On the contrary, consider generative sequential recommenders, which predict the upcoming item IDs based on input interaction history (a sequence of IDs). The differences between the sequential recommendation task and content tokenization could potentially be greater than those in text retrieval, potentially leading to differences in modeling. We expect to do more in-depth analysis in our future work.\n\n- Second, our proposed method also factors in the subsequent token as a residual of the previous one, assimilating this knowledge through the objective in Equation 4. \n\n- Third, we would like to note that the paper from Rajput et al. is still a concurrent work that appeared on ArXiv very close to the NeurIPS submission deadline. As a result, comparative analysis is not pursued yet within our submission paper.\n\n- Finally, we believe our approach is appropriate and effective for text retrieval tasks, and we will incorporate a more comprehensive discussion on Rajput et al.'s insightful study in our final version.\n\n**Unseen test set**\n\n- Thank you for your question. We did not split the test set because 84% of the queries in the MS MARCO test set are unseen, and all the queries in the BEIR test set are unseen [1]. Therefore, the results on MS MARCO and BEIR basically demonstrate the model's ability to retrieve on unseen. We will add clarification to our paper about this regard.\n\n**Why both contrastive loss and cross-entropy loss**\n\n- Since the proposed model jointly learns two objectives: decomposing/quantifying document semantic information into discrete docid, and generating relevant docid for a given query, both the contrastive learning loss for optimizing text representation, and the cross-entropy loss for optimizing generation, are employed. Due to the limited timeframe for a rebuttal, we would like to conduct a more in-depth study on the strengths and weaknesses between the joint-model and the pipeline-model in our future study. We thank you for these insightful comments.\n\n**W/o learning**\n\n- This variant directly employs the final docid assigned by the proposed system and trains a separate DSI model following the practice of DSI-QG. The model is trained without adopting progressive training or contrastive loss. This variant demonstrates that the generative retrieval model jointly trained with auto-encoding objectives can represent documents more sensibly based on semantics. This could potentially enhance performance on the less-optimized documents. Contrarily, the parameters obtained via cross-entropy loss on docid generation tasks might be less effective in conveying the semantic information of documents.\n\n**Qualitative analysis of docid**\n\n- Thanks for your suggestions. In our newly uploaded PDF, we add two qualitative analysis:\n\n- Figure 1 shows a case study of document content and the corresponding docid generated by GenRet. We can see that those documents with more similar docid have more relevant content. For example, we find that documents with docid started with 338-173 are related to Email, while documents with docid started with 338 are related to information exchange methods.\n\n- Figure 2 visualizes the codebook embedding and document embedding. We see that the codebook embedding does uniformly scatter in the document representation space, and meaningful clusters arise when documents are grouped by docids.\n\n- Figure 3 illustrates the word cloud of documents grouped by GenRet produced docid prefix. We can see that different positions of docid represent different levels of information, and the semantics within each cluster are related.\n\n\n**Commitment loss**\n\n- The commitment loss can be divided into two parts: the loss for previous docid token $z_{<t}$, and the loss for docid token at the current training step $z_t$. The former is critical for the model not to forget previous tokens during progressive training; the latter is important for the model to converge on the assignment of docid.\n\n- We have included detailed results of the ablation experiments in a newly uploaded PDF. Without commitment loss for previous tokens, the model's prediction accuracy for docid drops to less than 30% within 5 training epochs. Without commitment loss for the current token, the model's docid variance across epochs stays higher than 5%.\n\nWe hope our response can address your concerns.\n\n[1] Thakur, Nandan et al. BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models."}, {"review_id": "G8f0Xmz7c7", "reviewer": "Reviewer_xNtF", "summary": "As 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.", "questions": "**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?", "limitations": "As 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.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "1. 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.  ", "weaknesses": "1. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your time and valuable comments!\n\n**Presentation of method section**\n\n- Thanks for your suggestions, we will add more explanations to the final paper to improve presentation. \n\n- Firstly, the reconstruction loss optimizes both the encoder-decoder and the codebook. Concretely, $z_T$ is the codebook embedding that is closest in terms of inner-product distance to $d_T$. This computation involves a non-differentiable operation–argmax(·). For optimization of $d_T$ using $L_{rec}$, we employ straight-through gradient estimation [37] that allows for the transmission of the $z_T$ gradient to $d_T$. \n\n- Secondly, the parameters of both the encoder-decoder and codebook are shared between the tokenization model and the retrieval model. \n\n- Thirdly, at the inference stage, the query is entered into the encoder, and it is then processed by decoder+codebook in an autoregressive manner to output docid tokens.\n\n- Lastly, an intuitive explanation of re-assignment technique is that in instances where multiple documents receive the same ID, those with lower relative confidence are re-assigned. This re-assignment directs them towards less popular IDs. This process combined with the commitment loss; results in an increase in the prediction probability of these less frequented IDs, and an overall improvement in the diversity and uniformity of the assigned IDs.\n\n**Qualitative analysis**\n\n- Thanks for your suggestions. We visualize the codebook embedding and document embedding in the newly uploaded PDF files. We see that the codebook embedding does uniformly scatter in the document representation space.\n\n**Loss weighting**\n\n- Thanks for the question. During training, I add a factor of 0.1 to the reconstruction losses to balance the scale. We will add more clarification in our final paper.\n\n**Length of docid**\n\n- Thanks for the question. We list the final choice of hyper parameter M as follow, and we will add more detail our final paper.\n| NQ320K | MSMARCO320K | BEIR-Arg | BEIR-Covid | BEIR-NFC | BEIR-SciFact | BEIR-SciDocs | BEIR-FiQA |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| 3 | 4 | 2 | 4 | 2 | 2 | 3 | 3 |\n\n- A theoretical analysis suggests that given a certain number of documents $D$, the $M$ value is related to the diversity of ID assignments (e.g., the diversity metric defined on line 285). If the assignment is uniform, then $M \\propto log D$. Conversely, if the ID distribution is starkly imbalanced, then may $M \\propto D$. In practice, predicting the diversity of assignments is challenging. Observations indicate that the diversity of assignments decreases with an increase in docid length. Consequently, in our study the $M$ value is determined through empirical results, some of which are provided in the table above.\n\n**Limitation**\n\n- As we have discussed in conclusion, scaling the model to larger data typically requires longer docid length, more model capacity (parameters), more computation resources.In addition, we would like to add that: (i) our evaluation scale aligns with previous studies such as DSI and NCI. (ii) our proposed method outperforms existing generative retrieval systems on the unseen subset, and finally, application to larger datasets may pose additional problems e.g. increase in computational resources, which may fall outside the realm this study but could provide an interesting direction for future research.\n\nWe hope our response can address your concerns.\n\n[37] Oord, Aäron van den et al. Neural Discrete Representation Learning.\n"}]

Author rebuttal:
To all.

We appreciate all the reviewers for the constructive comments. We have included four figures in our newly updated PDF:

- Figure 1 presents a case study of document content matched with the relevant docid generated by GenRet. It suggests that documents with similar docids share closely related content. For instance, documents beginning with the docid 338-173 pertain to Email, while those starting with 338 concern various information exchange methods.

- Figure 2 shows a t-sne visualization of both codebook embedding and document embedding on NQ320K. The codebook embedding appears to distribute uniformly within the document representation space, producing meaningful clusters when documents are categorized by docids.

- Figure 3 shows a word cloud composed of documents grouped by GenRet-produced docid prefixes. It’s evident that documents of same group are semantically related.

- Figure 4 illuminates the results of an ablation experiment on commitment loss. The data suggests that commitment loss is beneficial for preventing the model from forgetting previous tokens and aids in enhancement of convergence.

We will add these results to our final paper. 

Best regard.


Meta-review: This is a well-written paper that proposes a straightforward but highly effective way of learning document identifiers for generative information retrieval (one of the most important problems in this area). The experimental results are strong, comprehensive, and convincing.

The reviewers raised a few concerns about this paper, but the most critical aspects of those were adequately addressed during the rebuttal.

Given that the strengths of this paper clearly outweigh its weaknesses, this paper is suitable for publication.

The authors are strongly encouraged to carefully consider all of the reviewer feedback, including the rebuttal-related discussion, and to take meaningful steps to incorporate it into the final version of their paper.
Decision: Accept (poster)