Learning to Tokenize for Generative Retrieval

Conventional document retrieval techniques are mainly based on the index-retrieve paradigm. It is challenging to optimize pipelines based on this paradigm in an end-to-end manner. As an alternative, generative retrieval represents documents as identifiers (docid) and retrieves documents by generating docids, enabling end-to-end modeling of document retrieval tasks. However, it is an open question how one should define the document identifiers. Current approaches to the task of defining document identifiers rely on fixed rule-based docids, such as the title of a document or the result of clustering BERT embeddings, which often fail to capture the complete semantic information of a document. We propose GenRet, a document tokenization learning method to address the challenge of defining document identifiers for generative retrieval. GenRet learns to tokenize documents into short discrete representations (i.e., docids) via a discrete auto-encoding approach. Three components are included in GenRet: (i) a tokenization model that produces docids for documents; (ii) a reconstruction model that learns to reconstruct a document based on a docid; and (iii) a sequence-to-sequence retrieval model that generates relevant document identifiers directly for a designated query. By using an auto-encoding framework, GenRet learns semantic docids in a fully end-to-end manner. We also develop a progressive training scheme to capture the autoregressive nature of docids and to stabilize training. We conduct experiments on the NQ320K, MS MARCO, and BEIR datasets to assess the effectiveness of GenRet. GenRet establishes the new state-of-the-art on the NQ320K dataset. Especially, compared to generative retrieval baselines, GenRet can achieve significant improvements on the unseen documents. GenRet also outperforms comparable baselines on MS MARCO and BEIR, demonstrating the method's generalizability.

Paper

References (60)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer xNtF6/10 · confidence 3/52023-07-04

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.

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. 2. 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. 3. 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. 2. 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.

Questions

**Q1:** Please take a look at **Weakness1** for suggestions to improve the presentation. **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? **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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

2 fair

Contribution

3 good

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.

Reviewer BkMu5/10 · confidence 4/52023-07-06

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.

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: * 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. * The evaluation using an unseen dataset is only done with the NQ dataset. * 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. [1] Shashank Rajput, Nikhil Mehta, Anima Singh, Raghunandan H Keshavan, Trung Vu, Lukasz Heldt, Lichan Hong, Yi Tay, Vinh Q Tran, Jonah Samost, et al. Recommender systems with generative retrieval. arXiv preprint arXiv:2305.05065, 2023.

Questions

Here are some questions I have: * 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. * Can you compare GENRET with [1]? * 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? * How important is the commitment loss? * 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. [1] Shashank Rajput, Nikhil Mehta, Anima Singh, Raghunandan H Keshavan, Trung Vu, Lukasz Heldt, Lichan Hong, Yi Tay, Vinh Q Tran, Jonah Samost, et al. Recommender systems with generative retrieval. arXiv preprint arXiv:2305.05065, 2023.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Yes, the authors to discuss limitations in terms of scale of dataset and generalization to different types/domains.

Reviewer BPS17/10 · confidence 3/52023-07-07

Summary

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

Strengths

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

Weaknesses

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

Questions

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

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Limitations were properly addressed.

Reviewer e88A6/10 · confidence 4/52023-07-07

Summary

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

Strengths

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

Weaknesses

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

Questions

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

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The BEIR results are still behind non-generative retrievers.

Reviewer zz1Z5/10 · confidence 5/52023-07-09

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. Experiments 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.

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. - Authors adopt diverse baselines - 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. For BEIR experiments, it is unclear how the six datasets are selected from all BEIR tasks. [1] Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval. Xiong et al. [2] https://cs.uwaterloo.ca/~jimmylin/publications/Nogueira_Lin_2019_docTTTTTquery-v2.pdf

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. - How was the six BEIR dataset picked?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

Would be nice to discuss how this method scales up to larger corpus, and whether / how it supports index updating.

Reviewer xNtF2023-08-13

No further problems for me

I have read the other reviews and the rebuttal. And I thank the authors for adequately addressing my concerns. For $\mathcal{L}_{rec}$, I now better understand how the encoder-decoder and codebook are optimized. The added qualitative results also seem to support the validity of the proposed method well. I'm willing to raise my **score** to "6" if the authors promise to include these results and analysis in the final submission.

Reviewer BkMu2023-08-21

As the authors say that they are willing to include comparisons to Rajput et. al. [1], I've raised my score.I would also recommend adding the qualitative analysis of docids.

Reviewer zz1Z2023-08-22

Thank you for the detailed response! They addressed many of my concerns and I'm willing to raise my rating. Please explain in the paper why document retrieval is focused instead of passage retrieval.

Authorsrebuttal2023-08-22

Thank you for your insightful feedback and for considering a higher score. A document in our context typically corresponds to a webpage, while a passage is a 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 NQ320K and BEIR. Passage retrieval, as a different task, could lead to different evaluation results. Thanks for your comments and we will incorporate an explanation in our final paper. We are also open to exploring other IR tasks, e.g., passage retrieval, in our future work. Best Regards.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC