Bootstrapped Pre-training with Dynamic Identifier Prediction for Generative Retrieval (2024)

Abstract: Generative retrieval uses differentiable search indexes to directly generate relevant document identifiers in response to a query. Recent studies have highlighted the potential of a strong generative retrieval model, trained with carefully crafted pre-training tasks, to enhance downstream retrieval tasks via fine-tuning. However, the full power of pre-training for generative retrieval remains underexploited due to its reliance on pre-defined static document identifiers, which may not align with evolving model parameters. In this work, we introduce BootRet, a bootstrapped pre-training method for generative retrieval that dynamically adjusts document identifiers during pre-training to accommodate the continuing memorization of the corpus. BootRet involves three key training phases: (i) initial identifier generation, (ii) pre-training via corpus indexing and relevance prediction tasks, and (iii) bootstrapping for identifier updates. To facilitate the pre-training phase, we further introduce noisy documents and pseudo-queries, generated by large language models, to resemble semantic connections in both indexing and retrieval tasks. Experimental results demonstrate that BootRet significantly outperforms existing pre-training generative retrieval baselines and performs well even in zero-shot settings.

## introduction document retrieval is an important task with widespread applications, such as question answering @xcite and fact verification @xcite , which aims to retrieve candidate documents from a huge document collection for a given query @xcite . currently, the dominant implementation is dense retrieval @xcite ), * corresponding author. using general language models, such as bart @xcite and t5 @xcite , as the base seq2seq model has become a popular choice in gr @xcite @xcite . on top of this, some work has designed pre-training objectives for gr. for example, @xcite proposed indexingand retrieval-based pre-training tasks; document pieces or pseudo-queries are used as input, and docids (e.g., product quantization code) are predicted as output with maximum likelihood estimation (mle). similarly, @xcite proposed retrieval-based tasks, which aim to construct and learn pairs of pseudo-queries and docids (i.e., wikipedia titles) from the corpus. these works demonstrate that applying specialized pre-trained models to gr yields superior results compared to using general language models. research challenges. while pre-training methods have shown their effectiveness, important limitations remain in the following: (i) the construction process of pre-defined docids is independent from the pre-training process. this results in a semantic gap between both processes, which could potentially hinder the retrieval performance. (ii) docids remain unchanged during pre-training. if the initial docids are not suitable, they cannot be further adjusted after the training begins. consequently, it may become challenging to learn semantics and relationships between documents, impeding the achievement of satisfactory retrieval performance. (iii) existing pre-training methods do not explicitly consider the interrelations between document-docid or query-docid pairs. the widely-used mle objective may result in difficulties in distinguishing among similar documents and docids. therefore, we argue that the model should enhance its discriminative and generalization ability. approach. to address these challenges, we introduce a general bootstrapped pre-training method for gr, called bootret. our objective is to dynamically adjust docids in accordance with the evolving model parameters during pre-training. the key idea is inspired by that the human brain updates the organization of existing knowledge to better match updated goals or contents in learning @xcite . bootret includes three key steps: (i) initial docid generation. we leverage the encoder of the initial model to encode documents and then obtain the product quantization code @xcite as the initial docids. we pre-train bootret based on two kinds of large scale text corpus, i.e., ms marco @xcite and wikipedia @xcite . we then fine-tune bootret on two representative downstream datasets widely used in gr research. the empirical experimental results show that bootret can achieve significant improvements over strong gr baselines. contributions. our main contributions are: (i) we propose a bootstrapped pre-training framework for gr to iteratively update the model parameters and docids. (ii) bootret demonstrates superior performance in downstream retrieval tasks. for instance, on the ms marco dataset, it outperforms the strong pre-training gr baseline, ultron @xcite , by 11.8% in terms of hits@1. (iii) additionally, bootret exhibits better zero-shot performance than other general language models. 2 related work generative retrieval. gr marks a new paradigm in document retrieval that generates identifier strings of documents as the retrieval target @xcite . the current design of docids can be categorized into two types. (i) pre-defined static docids. they remain unchanged during training, such as document titles @xcite , urls @xcite , product quantization code @xcite . this design is simple and shows decent performance @xcite , but the pre-defined process is independent of training. (ii) learnable docids. they are optimized jointly with the retrieval task @xcite . though these docids are dynamic, their optimization primarily targets retrieval. nevertheless, docids serve functions in both indexing and retrieval. in addition to widely-used supervised learning approaches @xcite @xcite , recent studies @xcite @xcite have explored pre-training for gr. however, each study adopts fixed docids, ignoring the potential mismatch between docids and the updated model. in contrast, our work dynamically updates both the docids and the evolving model to enhance the effectiveness. bootstrapping. the idea of bootstrapping training methods have garnered significant interest in various natural language processing tasks @xcite @xcite . the approach involves generating new training data or information based on the previous model to iteratively enhance its capabilities. while the techniques are sometimes used in conjunction with supervised learning @xcite , our scenario involves an unlabeled corpus without ideal docids. thus, we adopt a more unsupervised approach, iteratively refining the gr model and docids. dense retrieval. dense retrieval @xcite is currently the de facto solution for document retrieval. it focuses on representing documents and queries as dense vectors in continuous spaces, capturing semantic relationships. efficient vector search is facilitated by approximate nearest neighbor @xcite algorithms. further enhancements include using pre-trained models within a dual-encoder architecture @xcite and hard negative mining techniques @xcite . compared to dense retrieval, gr could achieve end-to-end global optimization. however, its current performance lags behind state-of-the-art methods in dense retrieval. ## method this section introduces the details of the bootret model proposed in this paper. as shown in figure 1 , (i) given a corpus @xmath0, . . . , d |d| }, we first construct an initial docid id i for each document d i in d. the initial docid set is denoted as i 0 d . we employ an encoder-decoder language model as the base model, where initial parameters are denoted as θ 0 . (ii) then, while keeping i 0 d unchanged, we carefully design two pre-training tasks. during pretraining, the model parameters are updated from θ 0 to θ 1 . (iii) subsequently, fixing θ 1 , we update the docids to i 1 d , thus completing one iteration. the updated docids can be used to further retrain the model for a next iteration. we define the t-th iteration as updating the model parameters from θ t-1 to θ t with fixed i ## experimental settings pre-training corpus. for pre-training we use two large, publicly available corpora: (i) english wikipedia, which contains tens of millions of well-written documents and we downloaded this dump @xcite for pre-training, and (ii) the ms marco document collection @xcite , which has about 3 million documents extracted from web documents using the bing search engine. for each corpus, we sample 500k documents and generated four noisy documents and five pseudo-queries, i.e., x, for each document. this results in 2.5m documents and 2.5m pseudo-queries for pre-training. bootret-bs wiki and bootret-bs ms denote the model pre-trained on wikipedia and ms marco, respectively. downstream retrieval datasets. we leverage two representative retrieval datasets. (i) ms marco document ranking dataset @xcite . following the setup of @xcite , we sample a subset of 300k documents for experimentation, denoted as ms 300k, containing 360k training queries, 6980 evaluation queries. these documents do not overlap with the ones used in pretraining. (ii) natural question (nq) @xcite has real questions and wikipedia documents, having about 228k documents with 307k training queries and 7.8k test queries. baselines. following typical gr research @xcite , we examine three types of baseline: (i) sparse retrieval baselines: bm25 @xcite , and doct5query @xcite . (ii) dense retrieval baselines: repbert @xcite , dpr @xcite , and ance @xcite . (iii) advanced gr baselines: dsi @xcite , @xcite , seal @xcite , dsi-qg @xcite , nci @xcite , ultron-pq @xcite , corpusbrain @xcite , gen-ret @xcite , and novo @xcite . for more details on our baselines, please refer to appendix a.1. evaluation metrics. following gr work @xcite , for nq, we use hit ratio (hits@k) with @xmath1, 10} as the metric. for ms 300k, we also use mean reciprocal rank (mrr@k) with @xmath2, 20} @xcite . implementation details. for pre-training, we use the llama-13b model @xcite to generate noisy documents and pseudo-queries. we initialize our model with t5-base (220m) @xcite . for docids, we set the length g of pq codes to 24, the number of clusters k to 256, and the dimension of vectors to 768. the hyperparameters for pre-training are set to α = β = λ = 1, γ = ρ = 2 and τ = 0.2. the batch size is 256. the adam optimizer with a learning rate of 5e-5 is used, and the sequence length of documents is set to 512. the max training step is 500k, with the first iteration occurring at step 100k, followed by iterations every 40k steps thereafter. for fine-tuning, we use the pre-trained model obtained from the last iteration to generate docids. models are further fine-tuned with document-docid pairs and labeled query-docid pairs with mle @xcite . following @xcite , we additionally generate 10 pseudoqueries for each document to enhance training. we set the learning rate as 1e-3, and the max training steps as 30k. other settings remain consistent with the pre-training stage. all models are trained on eight nvidia tesla a100 80gb gpus. for inference, we build a pre- ## experimental results this section presents the experimental findings. ## conclusion in this work, we proposed bootret, a bootstrapped pre-training method for gr, addressing the mismatch between pre-defined fixed docids and evolving model parameters in existing pre-training approaches. it dynamically adjusts docids based on the model pre-trained with two tasks. extensive experiments validate that bootret achieves superior performance compared to strong gr baselines on downstream tasks, even in the zero-shot setting.