SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hyper-parameters. In our experiments, we utilize the T5-Base model [27] as the base Transformer and initialize a new codebook embedding  $\mathbf{E}_t$  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 on the number of documents present. For datasets containing a larger number of candidate documents, a larger value of  $M$  is set to ensure that all documents are assigned unique document ids. In the docid re-assignment, the hyper-parameter  $\epsilon$  is set to 1.0, and the Sinkhorn-Knopp algorithm is executed for 100 iterations.
Indexing with query generation. Following previous work [47, 41, 40], we use query generation models to generate synthetic (query, document) pairs for data augmentation. Specifically, we use the pre-trained query generation model from DocT5Query [24] to augment the NQ and MS MARCO datasets. In query generation, we use nucleus sampling with parameters  $p = 0.8$ ,  $t = 0.8$  and generate five queries for each document in the collection. For the BEIR datasets, we use the queries generated by GPL [40]. GPL uses a DocT5Query [24] generator trained on MS MARCO to generate about 250K queries for each BEIR dataset. Note that the query generator used for BEIR is purely trained on MS MARCO (without using any training data of BEIR) and thus conforms to the zero-shot setting of BEIR [38, 40].
Training and inference. The proposed models and the reproduced baselines are implemented with PyTorch 1.7.1 and HuggingFace transformers 4.22.2. We optimize the model using AdamW and set the learning rate to  $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 a factor of 0.1 to the reconstruction losses to balance the scale. In progressive training, we first warm up the model for 5K steps and then initialize the codebook using the clustering centroids as mentioned in Section 3.2.1. We use constrained clustering<sup>4</sup> to obtain diverse clustering results. During inference, we use beam search with constrained decoding [5] and a beam size of 100.
<sup>4</sup>https://github.com/joshlk/k-means-constrained
Table 1: Results on Natural Questions (NQ320K). The results of the methods marked with  $\dagger$  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.  $\sharp$  and  $\sharp$  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 are indicated in boldface.
Method	Full test	Seen test	Unseen test
R@1	R@10	R@100	MRR	R@1	R@10	R@100	MRR	R@1	R@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
DocT5Query [24]	38.0	69.3	86.1	48.9	35.1	68.3	86.4	46.7	48.5	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-T5†[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
GTR-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	83.6	90.5	74.7	1.3	7.2	31.5	3.5
SEAL [1]	59.9	81.2	90.9	67.7	-	-	-	-	-	-	-	-
CGR-Contra [17]	63.4	81.1	-	-	-	-	-	-	-	-	-	-
DSI-QG†[47]	63.1	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**
Results on NQ320K. In Table 1, we list the results on NQ320K. GENRET outperforms both the strong pre-trained dense retrieval model, GTR, and the previous best generative retrieval method, NCI, thereby establishing a new state-of-the-art on the NQ320K dataset. Furthermore, our results reveal that existing generative retrieval methods perform well on the seen test but lag behind dense retrieval methods on the unseen test. For example, NCI obtains an MRR@100 of 76.8 on the seen test, which is higher than the MRR@100 of 65.3 obtained by GTR-Base. However, on unseen test data, NCI performs worse than GTR-Base. In contrast, GENRET performs well on both seen and unseen test data. This result highlights the ability of GENRET to combine the advantages of both dense and generative retrieval by learning discrete docids with semantics through end-to-end optimization.
Results on MS MARCO. Table 2 presents the results on the MS MARCO dataset. GENRET outperforms generative retrieval methods such as Ultron and dense retrieval baselines such as ANCE and Sentence-T5. Furthermore, previous generative retrieval methods (e.g., GENRE, Ultron) utilizing metadata such as the title and URL, while exhibiting decent performance on the NQ320K dataset, underperform in comparison to dense retrieval and sparse retrieval methods on the MS MARCO dataset. This may be because the NQ320K dataset retrieves Wikipedia documents, where metadata like the title effectively capture the semantics of the document. In the case of the MS MARCO dataset, which is a web search dataset, the metadata often does not adequately characterize the documents, resulting in a decline in performance of the generative retrieval model. In contrast, GENRET learns to generate semantic docids that effectively enhance the generative retrieval model.
Results on BEIR. Table 3 lists the results of the baselines and GENRET on six datasets of BEIR. These datasets represent a diverse range of information retrieval scenarios. On average, GENRET outperforms strong baselines including BM25 and ST5 GPL, and achieves competitive results compared to previous-best sparse and dense retrieval methods. Additionally, GENRET demonstrates a significant improvement over the previous generative retrieval model GENRE that utilizes titles as docids. Furthermore, GENRE performs poorly on some datasets, such as BEIR-Covid and BEIRSciDocs. This may be because the titles of the documents in these datasets do not adequately capture their semantic content.
Table 2: Results on MS MARCO. The results of the methods marked with  $\dagger$  are from our own re-implementation, other results are cited from the original paper or implemented using official code. The best results are indicated in boldface.
Method	R@1	R@10	R@100	MRR
Sparse retrieval
BM25 [32]	39.1	69.1	86.2	48.6
DocT5Query [24]	46.7	76.5	90.4	56.2
Dense retrieval
ANCE [42]	45.6	75.7	89.6	55.6
Sentence-T5† [22]	41.8	75.4	91.2	52.8
Generative retrieval
GENRE† [5]	35.6	57.6	79.1	42.3
Ultron-URL [46]	29.6	67.8	-	40.0
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
Table 3: nDCG@10 results on BEIR. The results of the methods marked with  $\dagger$  are from our own re-implementation, other results are cited from the original paper or implemented using official code. ST5 GPL denotes Sentence-T5 trained on GPL datasets [40].
Method	Arg	Covid	NFC	SciF	SciD	FiQA	Avg.
Sparse retrieval
BM25 [32]	29.1	58.9	33.5	67.4	14.8	23.6	37.8
DocT5Query [24]	34.9	71.3	32.8	67.5	16.2	29.1	41.9
Dense retrieval
ST5 GPL†[22]	32.1	74.4	30.1	58.6	12.7	26.0	39.0
Contriever [11]	40.0	68.8	33.5	61.4	16.3	30.7	41.8
Generative retrieval
GENRE†[47]	42.5	14.7	20.0	42.3	6.8	11.6	30.0
Ours	34.3	71.8	31.6	63.9	14.9	30.2	41.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 a box with a docid and a diversity metric  $d$ , which is calculated by:  $d = 1 - \frac{1}{2n}\sum_{j=1}^{K}|n_j - n_u|$ , where  $|\cdot|$  represents the absolute value,  $n$  denotes the total number of documents,  $n_j$  denotes the number of documents that have a docid  $= j$ , and  $n_u = \frac{n}{K}$  is the expected number of documents per docid under the uniform distribution.
The results demonstrate the superiority of GENRET (represented by the yellow line) in terms of distribution uniformity. It uses all the potential docid  $k = 512$  and achieves the highest diversity metric with a value of  $d = 0.90$ . The method without docid reassignment also yields a relatively balanced distribution, with a diversity metric of  $d = 0.77$ . However, the distribution of the method without diverse codebook initialization is highly uneven, which could be due to the fact that most of the randomly initialized codebook embeddings are not selected by the model during the initial training phase, leading to a lack of update and further selection in subsequent training. Additionally, the models without diverse clustering tend to converge to a trivial solution where all documents are assigned the same docid.
In Figure 3 (right), the results of two ablated variants are presented. First, GENRET w/o learning is a generative model that has been trained directly using the final output docid from GENRET, without utilizing the proposed learning scheme. Its retrieval performance is comparable to that of GENRET on seen test data; however, it is significantly lower on unseen test data. This variant demonstrates that the generative retrieval model jointly trained with auto-encoding objectives can represent documents more sensibly based on semantics. This could enhance performance on the less-optimized documents. Contrarily, the parameters obtained via cross-entropy loss on docid generation tasks are less effective in conveying the semantic information of documents. Secondly, GENRET w/ T5-Small uses a small model, and its performance is inferior to that of GENRET using T5-Base. However, the gap between the performance on seen and unseen test data is smaller, which could be attributed to the limited fitting capacity of the small model.
In Appendix C, we evaluate the proposed model's capacity to retrieve new documents and find that it performs well in adapting to new documents compared to existing document tokenization approaches. Additionally, in Appendix D, we analyze the efficiency of various retrieval models in comparison to the various baselines in terms of memory usage, offline, and online latency, and show the advantages of the proposed model.
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 a word cloud of documents grouped by docid prefixes. It is evident that documents within the same group are semantically related. For example, major words for documents with the docid prefix 338 are mail and stamp. When a second-level docid 173 is added, the corresponding documents become more specifically related to email. With the addition of a third-level docid 1, the document group becomes specifically associated with email marketing (docid: 338-173-1). Similar patterns can be observed in the other three cases. The case study shows that there is a hierarchical semantic structure within the learned docids.
Figure 5 in Appendix E visualizes the codebook embedding and document embedding. We see that the codebook embedding appears to distribute uniformly within the document representation space, producing meaningful clusters when documents are categorized by docids. We also find that the uniformity of the embedding distribution decreases with increasing docid-length. This may be due to the fact that uniform segmentation is more challenging as the semantic granularity becomes finer.
Figure 4: Left: Document titles along with their corresponding docids. It is observed that documents with similar docids tend to 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.
Sparse and dense retrieval. Traditional sparse retrieval calculates the document score using term matching metrics such as TF-IDF [31], query likelihood [16], or BM25 [32]. Sparse retrieval is widely used in practice due to its efficiency, but often suffers from the lexical mismatches [18]. Dense
retrieval (DR) addresses this by presenting queries and documents in dense vectors and calculating their similarities with the inner product or cosine similarity [13]. Various techniques have been proposed to improve DR models, such as hard negative mining [42, 26], late interaction [14, 34], knowledge distillation [10, 19], and pre-training [30, 23, 11]. Despite their success, DR approaches have several limitations [5, 21]: (i) DR models employ an index-retrieval pipeline with a fixed search procedure (MIPS), making it difficult to optimize the model end-to-end [37, 41]. (ii) Training DR models relies on contrastive learning [13] to distinguish positives from negatives, which is inconsistent with large LMs training objectives [3] and fails to fully utilize the capabilities of pretrained LMs [1, 35].
Generative retrieval. Generative retrieval is gaining attention. It retrieves documents by generating their docid using a generative model like T5. Generative retrieval presents an end-to-end solution for document retrieval tasks [37, 21] and allows for better exploitation of the capabilities of large generative LMs [1]. Cao et al. [5] first propose an autoregressive entity retrieval model to retrieve documents by generating titles. Tay et al. [37] propose a differentiable search index (DSI) and represent the document as atomic id, naive string, or semantic string. Bevilacqua et al. [1] suggest using arbitrary spans of a document as docids. Additionally, multiple-stage pre-training [7, 46], query generation [41, 47, 46], contextualized embedding [17], and continual learning [20], have been explored in recent studies. Recently, Tang et al. [36] introduce a query-based docid and the rehearsal-based document indexing to improve DSI. However, existing generative retrieval models have a limitation in that they rely on fixed document tokenization to produce docids, which often fails to capture the semantic information of a document [37]. It is an open question how one should define the docids. To further capture document semantics in the docid, we propose document tokenization learning methods. The semantic docid is generated by the proposed discrete auto-encoding learning scheme in an end-to-end manner. Concurrently, Rajput et al. [28] propose a RQ-VAE module to produce semantic docids for generative recommender systems. As a comparison, our proposed method jointly models the tokenization and retrieval tasks with shared parameters to better align the model's representation of the two tasks and are optimized in an end-to-end manner.
Discrete representation learning. Learning discrete representations using neural networks is an important research area in machine learning. For images, Rolfe [33] proposes the discrete variational autoencoder, and VQ-VAE [39] learns quantized representations via vector quantization. Dalle [29] uses an autoregressive model to generate discrete image representation for text-to-image generation. Recently, representation learning has attracted considerable attention in NLP tasks, for tasks such as machine translation [48], dialogue generation [45], and text classification [12, 43]. For document retrieval, RepCONC [44] uses a discrete representation learning method based on constrained clustering for vector compression. We propose a document tokenization learning method for generative retrieval, which captures the autoregressive nature of docids by progressive training and enhances the diversity of docids by diverse clustering techniques.
This paper has proposed a document tokenization learning method for generative retrieval, named GENRET. The proposed method learns to tokenize documents into short discrete representations (i.e., docids) via a discrete auto-encoding approach, which ensures the semantics of the generated docids. A progressive training method and two diverse clustering techniques have been proposed to enhance the model's training. Empirical results on various document retrieval datasets have demonstrated the effectiveness of the proposed method. Especially, GENRET achieves outperformance on unseen documents and can be well generalized to multiple retrieval tasks.
The limitations of this work include experiments only on moderately sized datasets like NQ320K. The employed data are in sufficient quantity to validate the effectiveness of the proposed method, but application to larger-scale data may require more model parameters and computational resources. Another aspect for improvement is the generalization of the model to unoptimized document collections in different types or domains, as compared to continuous embedding approaches. We recognize this as an important research question for generative retrieval but believe that addressing this is beyond the scope of this paper. In future work, we would like to extend the approach to large document collections. We also plan to explore generative pre-training for document tokenization using large-scale language models. Additionally, we intend to investigate the dynamic adaptation of docid prefixes for progressive training.
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 (2019JZZY010129), the Fundamental Research Funds of Shandong University, the China Scholarship Council under grant nr. 202206220085, the Hybrid Intelligence Center, a 10-year program funded by the Dutch Ministry of Education, Culture and Science through the Netherlands Organisation for Scientific Research, https://hybrid-intelligence-centre.nl, and project LESSEN with project number NWA.1389.20.183 of the research program NWA ORC 2020/21, which is (partly) financed by the Dutch Research Council (NWO). All content represents the opinion of the authors, which is not necessarily shared or endorsed by their respective employers and/or sponsors.
[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, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, T. J. Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In NeurIPS 2020, 2020.
[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 CComprehension Dataset. ArXiv, abs/1611.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, Yiqun Liu, Yixing Fan, and Xueqi Cheng. Corpus-Brain: Pre-train a generative retrieval model for knowledge-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 in continuous space. ArXiv, abs/1811.08008, 2018.
[10] Sebastian Hofstätter, Sheng-Chieh Lin, Jheng-Hong Yang, Jimmy J. Lin, and Allan Hanbury. Efficiently teaching an effective dense retriever with balanced topic aware sampling. In SIGIR 2021, 2021.
[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 representations for low-resource text classification. In ACL 2020, 2020.
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	# Docs	# Test queries	# Train pairs
NQ320K	109,739	7,830 / 6,075 / 1,755	307,373
MS MARCO	323,569	5,187 / 807 / 4,380	366,235
BEIR-Arg	8,674	1,406	-
BEIR-Covid	171,332	50	-
BEIR-NFC	3,633	323	-
BEIR-SciFact	5,183	300	-
BEIR-SciDocs	25,657	1,000	-
BEIR-FiQA	57,638	648	-
We conduct experiments on three document<sup>5</sup> retrieval datasets: NQ320K, MS MARCO, and BEIR.
NQ320K. NQ320K is a popular dataset for evaluating generative retrieval models [37, 41]. It is based on the Natural Questions (NQ) dataset proposed by Google [15]. NQ320K consists of 320k query-document pairs, where the documents are gathered from Wikipedia pages, and the queries are natural language questions. We follow the evaluation setup in NCI [41] and further split the test set into two subsets: seen test, in which the annotated target documents of the queries are included in the training set; and unseen test, in which no labeled document is included in the training set.
Note that the NQ320K dataset we utilized has been pre-processed based on the NCI [41] and includes approximately 100k documents. This differs from the DSI approach [37], which processed a version of the NQ320K dataset containing about 200k documents. The distinction lies in the method used to remove duplicate documents. In our case, we eliminated duplicates by comparing document titles, whereas DSI employed URLs. For example, pages https://en.wikipedia.org//w/index.php?title=Statue_of_Liberty&amp;oldid=804877528 and https://en.wikipedia.org/w/index.php?title=Statue_of_Liberty&amp;oldid=834310497 are two versions of entity "Statue of Liberty". DSI NQ320K treats them as separate documents, whereas our implementation considers them as a single document. We have found that the content of different versions of the same entity's pages is usually almost identical, with only minor variations, often occurring in later parts of the document. Consequently, we believe that distinguishing between different versions of an entity is beyond the model's capabilities, and that using different versions of the same entity as negative examples in training may hurt model performance.