Source-linked AI summary
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
TL;DR
Document retrieval pipelines are difficult to optimize end to end, and how to define docids that represent document semantics remains an open question. GenRet learns discrete docids through auto-encoding and uses them for generative retrieval. It establishes a new state of the art on NQ320K and reports gains on unseen documents and other retrieval benchmarks.
Problem
How to define docids that capture document semantics remains open, and existing rule-based tokenizers may generalize poorly to new or out-of-distribution documents.
Source: §1 · pp. 1–2Method
GenRet jointly trains document tokenization, reconstruction, and query-to-docid generation in a discrete auto-encoding framework.
Source: §1 · p. 2Results
Across NQ320K, MS MARCO, and six BEIR datasets, GenRet outperforms generative retrieval baselines; its NQ320K unseen-test improvement is +14% relative to the best baseline.
Source: §1 · p. 2Takeaways & Limitations
GenRet's results indicate effectiveness on unseen documents and generalization across multiple retrieval tasks.
Source: §7 · p. 9Duplicate docids can result, in which case the corresponding documents are retrieved in an arbitrary order.
Source: §5.2 · p. 7
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The introduction motivates generative retrieval as an end-to-end alternative to fixed index-retrieve pipelines, but leaves docid design as an open problem. GenRet learns semantic docids with discrete auto-encoding and reports gains, especially on unseen documents.
1 INTRODUCTION
- Dense retrieval's fixed MIPS search procedure makes end-to-end joint optimization difficult, while generative retrieval generates docids autoregressively for queries.The introduction also notes that dense-retrieval learning strategies can mismatch pre-training objectives such as next-token prediction.Source: §1 · p. 1
- Defining docids remains open: rule-based titles, URLs, or embedding clusters are ad hoc and may generalize poorly to new or out-of-distribution documents.Document tokenization defines how documents are distributed in semantic space.Source: §1 · pp. 1–2
- GenRet jointly optimizes document tokenization, reconstruction, and query-to-docid retrieval in a discrete auto-encoding framework to learn semantic docids.The paper presents it as a tokenization-learning method for document retrieval.Source: §1 · p. 2
- Progressive training stabilizes autoregressive docid learning by fixing optimized prefixes and using reconstruction, commitment, and retrieval losses.These losses respectively target document reconstruction, docid commitment to avoid forgetting, and end-to-end retrieval performance.Source: §1 · p. 2
- Parameter initialization and docid reassignment using diverse clustering techniques are proposed to increase the diversity of generated docids.The paper identifies docid diversity as a challenge for auto-encoding-based optimization.Source: §1 · p. 2
- GenRet establishes NQ320K state of the art with +14% relative improvement on the unseen test set versus the best baseline, and outperforms generative methods on MS MARCO and six BEIR datasets.It also achieves competitive results compared with the best dense retrieval model on MS MARCO and BEIR.Source: §1 · p. 2
2 PRELIMINARIES
The preliminaries formalize document retrieval as finding a relevant document for a query and describe docids as shorter discrete document representations. Generative retrieval retrieves by generating those identifiers autoregressively, with constrained decoding used to keep outputs valid.
2 PRELIMINARIES
- Document retrieval is framed as finding a relevant document in collection D for a search query.Documents are represented as sequences of tokens.Source: §2 · p. 2
- A docid is a shorter fixed-length sequence of discrete tokens intended to distinguish documents and capture their semantics, simplifying training and inference.Each position is modeled as a K-way categorical variable.Source: §2 · p. 2
- After documents receive docids, the retrieval model generates a query's relevant docids autoregressively and uses constrained decoding to restrict outputs to identifiers in the corpus.A prefix tree represents valid docids, invalid outputs receive zero probability, and beam search retrieves multiple documents.Source: §2 · p. 2
3 GENRET
GenRet learns discrete semantic docids through an auto-encoding framework: documents are tokenized, reconstructed from their docids, and retrieved by generating docids. Its training addresses autoregressive docid learning and imbalanced assignments through progressive optimization and diversity techniques.
3 GENRET
- GenRet tokenizes documents into discrete docids that can reconstruct the original document, then retrieves by generating relevant docids.The overview depicts tokenization, reconstruction, and autoregressive generation as the method's pipeline.Source: §3 · p. 2, Figure 2 · p. 3
3.1 Document tokenization and retrieval model
- At each timestep, the tokenization model uses a codebook of K embeddings to select a discrete token for the document.Each codebook embedding represents a segmentation; the token probability is calculated from the document representation and codebook.Source: §3.1 · p. 3
3.2 Document reconstruction model
- The reconstruction model retrieves the target document from documents sharing its docid prefix, learning residual semantics not captured by earlier tokens.Its input is the docid representation, and its output is the associated document.Source: §3.2 · pp. 3–4
3.3 Model optimization
- GenRet uses a progressive auto-encoding scheme to address challenges in autoregressive docid learning and docid diversity.The challenges include changing, under-optimized prefixes and imbalanced assignment in which a few docids are used for many documents.Source: §3.3 · p. 4
3.3.1 Progressive training scheme.
- Progressive training optimizes the docid at each step while keeping earlier docids and other parameters fixed.At step T, the corresponding docid and its generation parameters are updated; previously produced docids remain fixed.Source: §3.3.1 · p. 4, Figure 3 · p. 4
3.3.2 Diverse clustering technique.
- Codebook initialization clusters document representations with constrained K-means, while batch-level docid reassignment encourages distinct assignments.The initialization seeks balanced semantic-space segmentation, and reassignment seeks balanced docid assignments.Source: §3.3.2 · pp. 4–5
4 EXPERIMENTAL SETUP
The experiments cover NQ320K, MS MARCO, and six heterogeneous BEIR tasks, with dataset-specific evaluation metrics. Baselines span sparse, dense, and generative retrieval, and experiments use a shared T5-Base setup with specified training and inference procedures.
4.1 Datasets
- NQ320K uses 320k query-document pairs and separates test queries into seen and unseen subsets; MS MARCO evaluates queries against a sampled document subset, while six BEIR tasks cover heterogeneous collections.All BEIR test queries are unseen.Source: §4.1 · p. 5, Table 1 · p. 5
4.2 Evaluation metrics
- NQ320K is evaluated with Recall@{1,10,100} and MRR@100, MS MARCO with Recall@{1, 10, 100} and MRR@10, and BEIR with average nDCG@10.BEIR averages nDCG@10 across its downstream datasets.Source: §4.2 · p. 5
4.3 Baselines
- Baselines include BM25 and DocT5Query for sparse retrieval, and DPR, ANCE, Sentence-T5, GTR, and Contriever for dense retrieval.The dense methods use dual encoders; GTR pre-trains Sentence-T5 on billions of paired data using contrastive learning.Source: §4.3 · p. 5
- Generative baselines use titles, clustering IDs, n-grams, query augmentation, or multiple identifier types; Sentence-T5, GENRE, DSI-QG, and GenRet provide a T5-based comparison with differing outputs.NCI augments training data with 15 generated queries per document; Ultron uses URL, PQ, and Atomic identifiers.Source: §4.3 · pp. 5–6
4.4 Implementation details
- Experiments use T5-Base, 512 clusters, dataset-dependent docid lengths, and Sinkhorn-Knopp reassignment; indexing augments NQ and MS MARCO with five generated queries per document.BEIR uses GPL-generated queries; training uses AdamW and inference uses constrained beam search with beam size 100.Source: §4.4 · p. 6
5 EXPERIMENTAL RESULTS
GenRet establishes the NQ320K state of the art and performs strongly on MS MARCO and BEIR. Further experiments examine unseen-document retrieval, docid distributions, efficiency, and a retrieval example.
5.1 Main results
- GenRet sets NQ320K’s state of the art, beats prior generative retrieval on MS MARCO while matching GTR, and averages above BM25 and GTR-Base on BEIR.On BEIR, it also achieves better results than ST5 GPL and GENRE, and is competitive with state-of-the-art sparse and dense methods.Source: §5.1 · p. 6, Table 2 · p. 7, Table 3 · p. 7
5.2 Performance on retrieving new documents
- nDCG@10 12.3 on BEIR-SciDocs was comparable to target-trained models after training on NQ and evaluating on unseen NQ320K and BEIR collections.The tokenization model generates docids for target-collection documents; duplicate docids can cause corresponding documents to be returned in arbitrary order.Source: §5.2 · pp. 6–7, Table 5 · p. 8
5.3 Analytical experiments
- GenRet uses all 512 potential docids and reaches diversity d = 0.90; without docid reassignment, d = 0.77, while removing diverse clustering yields a trivial shared docid.Without diverse codebook initialization, the docid distribution is highly uneven because many randomly initialized codebook embeddings are not selected and updated early in training.Source: §5.3 · p. 8, Figure 4 · p. 8
5.4 Efficiency analysis
- Generative models use less memory than dense and sparse methods, but GenRet requires longer offline indexing; its shorter docids improve online speed relative to DSI and GENRE.Generative-model memory depends on model parameters, whereas dense and sparse methods also store document embeddings; generative retrieval latency varies with beam size and docid length.Source: Table 6 · p. 8, §5.4 · p. 8
5.3 Analytical experiments
- Without progressive auto-encoding, retrieval is comparable on seen data but significantly lower on unseen data; T5-Small performs worse overall while narrowing the seen–unseen gap.The authors attribute the smaller gap to the smaller model’s limited fitting capacity.Source: §5.3 · p. 8, Figure 4 · p. 8
5.5 Case study
- For the NQ320K query, GenRet returns the target docid unlike the baselines, attending to topic-related words at t = 1 and country-related words at t = 2.The example highlights “Appellate” at t = 1 and “United States” at t = 2.Source: §5.5 · pp. 8–9, Table 7 · p. 9
6 RELATED WORK
Related work reviews sparse and dense retrieval alongside generative retrieval, whose fixed document tokenization motivates GenRet’s learned semantic docids. It also situates GenRet among discrete representation methods.
Sparse retrieval
- Sparse retrieval scores documents through term matching, including TF-IDF, query likelihood, and BM25, but can suffer from lexical mismatch.Source: §6 · p. 9
Dense retrieval
Generative retrieval
- Generative retrieval generates docids with a generative model, but prior approaches rely on fixed tokenization that may fail to capture document semantics [1] [5] [28] [42].Prior docids include titles, atomic or semantic strings, arbitrary document spans, and other fixed schemes [1] [5] [42].Source: §6 · p. 9
7 CONCLUSIONS
GenRet learns short discrete document identifiers with a discrete auto-encoding approach, supported by progressive training and diverse clustering techniques. Experiments show effectiveness across retrieval datasets, especially on unseen documents.
7 CONCLUSIONS
- GenRet learns short discrete document representations, or docids, through discrete auto-encoding to preserve their semantics.Source: §7 · p. 9
- Progressive training and two diverse clustering techniques are proposed to enhance model training.Source: §7 · p. 9
- Experiments across retrieval datasets show effectiveness, with outperformance on unseen documents and generalization to multiple retrieval tasks.Source: §7 · p. 9
- Future work will extend the approach to large collections, explore generative pre-training, and investigate dynamic adaptation of docid prefixes.Source: §7 · p. 9