Unsupervised Corpus Aware Language Model Pre-training for Dense Passage Retrieval (2022)

Abstract: Recent research demonstrates the effectiveness of using fine-tuned language models (LM) for dense retrieval. However, dense retrievers are hard to train, typically requiring heavily engineered fine-tuning pipelines to realize their full potential. In this paper, we identify and address two underlying problems of dense retrievers: i) fragility to training data noise and ii) requiring large batches to robustly learn the embedding space. We use the recently proposed Condenser pre-training architecture, which learns to condense information into the dense vector through LM pre-training. On top of it, we propose coCondenser, which adds an unsupervised corpus-level contrastive loss to warm up the passage embedding space. Experiments on MS-MARCO, Natural Question, and Trivia QA datasets show that coCondenser removes the need for heavy data engineering such as augmentation, synthesis, or filtering, and the need for large batch training. It shows comparable performance to RocketQA, a state-of-the-art, heavily engineered system, using simple small batch fine-tuning.

## introduction building upon the advancements of pre-trained language models (lm; @xcite ; @xcite al. (2019) ), dense retrieval has become an effective paradigm for text retrieval @xcite @xcite @xcite . recent research has however found that fine-tuning dense retrievers to realize their capacity requires carefully designed fine-tuning techniques. early works include iterative negative mining @xcite and multi-vector representations @xcite . the recent rocketqa system @xcite significantly improves the performance of a dense retriever by designing an optimized fine-tuning pipeline that includes i) denoising hard negatives, which corrects mislabeling, and ii) large batch training. while this is very effective, the entire pipeline is very heavy in computation and not feasible for people who do not have tremendous hardware resources, like those in academia. in this paper, we ask, instead of directly using the pipeline, can we take the insights of rocketqa to perform language model pre-training such that the pre-trained model can be easily fine-tuned on any target query set. concretely, we ask what the optimized training in rocketqa solves. we hypothesize that typical lms are sensitive to mislabeling, which can cause detrimental updates to the model weights. denoising can effectively remove the bad samples and their updates. on the other hand, for most lms, the cls vectors are either trained with a simple task @xcite or not explicitly trained at all @xcite . these vectors are far from being able to form an embedding space of passages @xcite . the large training batches in rocketqa help the lm to stably learn to form the full embedding space. to this end, we want to pre-train an lm such that it is locally noise-resistant and has a well-structured global embedding space. for noise resistance, we borrow the condenser pre-training architecture @xcite , which performs language model pretraining actively conditioning on the cls vector. it produces an information-rich cls representation that can robustly condense an input sequence. we then introduce a simple corpus level contrastive learning objective: given a target corpus of documents to retrieve from, at each training step sample text span pairs from a batch of documents and train the model such that the cls embeddings of two spans from the same document are close and spans from different documents are far apart. combining the two, we propose cocondenser pre-training, which unsupervisedly learns a corpus-aware pretrained model for dense retrieval. in this paper, we test cocondenser pre-training on two popular corpora, wikipedia and ms-marco. both have served as information sources for a wide range of tasks. this popularity justifies pre-training models specifically for each of them. we directly fine-tune the pre-trained cocondenser using small training batches without data engineering. on natural question, triviaqa, and ms-marco passage ranking tasks, we found that the resulting models perform on-par or better than rocketqa and other contemporary methods. ## related work dense retrieval transformer lm has advanced the state-of-the-art of many nlp tasks @xcite @xcite including dense retrieval. @xcite are among the first to demonstrate the effectiveness of transformer dense retrievers. they proposed a simple inverse cloze task (ict) method to further pre-train bert @xcite . follow-up works explored other pretraining tasks @xcite as well end-toend co-training of reader and retriever @xcite . @xcite is the first to discover that careful fine-tuning can learn effective dense retriever directly from bert. later works then started to investigate ways to further improve fine-tuning @xcite . among them, @xcite proposed the rock-etqa fine-tuning pipeline which hugely advanced the performance of dense retrievers. until very recently, pre-training for dense retrieval has been left unexplored. a concurrent work, dpr-paq @xcite , revisits pretraining and proposes domain matched pre-training, using a 65-million-size synthetic qa pair dataset generated with pre-trained natural question and trivia qa pipelines to pre-train dense retrievers. this paper uses a recently proposed dense retrieval pre-training architecture, condenser @xcite . unlike previous works that design pre-training tasks, condenser explored the idea of designing a special pre-training architecture to improve representation effectiveness. one reason why dense retrieval is of immediate great value is that there is a rich literature that studies efficient dense retrieval for first stage retrieval @xcite . there are also mature dense retrieval libraries, such as faiss @xcite . by pre-encoding the corpus into a mips index, retrieval can run online with millisecond-level latency @xcite . contrastive learning contrastive learning has become a very popular topic in computer vision @xcite . recent works have brought the idea to natural language processing to learn high-quality sentence representation @xcite . in this work, we use contrastive learning to do pretraining for dense retrieval. different from earlier work, instead of individual representations (giorgi et al., 2020), we are interested in the full learned embedding space, which we will use to warm start the retriever. the large batch requirement had been a limiting factor in contrastive learning @xcite under resource-limited setups where gpu (accelerator) memory is not sufficiently large. in general, this extends to any training procedure that uses contrastive loss, including dense retrieval pretraining @xcite . @xcite recently devised a gradient cache technique that upper-bounds peak memory usage of contrastive learning to almost constant. in subsection 3.3, we show how to adapt it for cocondenser pre-training. ## method in this section, we first give a brief review of condenser. then we discuss how to extend it to co-condenser and how to perform memory-efficient cocondenser pre-training. oven [mask] apple pie [cls] oven [mask] apple pie [cls] oven [mask] apple pie [cls] oven [mask] apple pie [cls] oven [mask] apple pie [cls] head (pre-train only) oven [mask] apple pie [cls] oven [mask] apple [cls] pie late early figure 1: condenser: shown are 2 early and 2 late backbone layers. in our experiments each have 6 layers. condenser head is dropped during fine-tuning. this illustration is taken from the condenser paper @xcite . ## experiments in this section, we first describe the implementation details of cocondenser pre-training. we then conduct dense retrieval experiments to test the effectiveness of fine-tuned cocondenser retrievers. ## analysis recall the two desired properties of cocondernser are local noise resistance to mislabeling and a wellstructured, pre-trained embedding space. in this section, to investigate the former, we introduce and compare with a knowledge distillation @xcite setup where we substitute noisy hard labels with soft labels from a cross-encoder. for the latter, we measure the quality of 1st stage retriever mined negatives to see if cocondenser embedding can help find related but not relevant hard negatives. we also provide ablation studies of loss components and, pre-training and fine-tuning stages. ## conclusion this paper introduces cocondenser, an unsupervised corpus-aware language model pre-training method. we demonstrate proper pre-training can establish not only language understanding ability but also corpus-level representation power. leveraging the condenser architecture and a corpus aware contrastive loss, cocondenser acquires two important properties for dense retrieval, noise resistance, and structured embedding space. this corpus-aware pre-training needs to be done once for a search corpus and is query agnostic. the learned model can be shared among various types of end task queries. experiments show that cocondenser can drastically reduce the costs of fine-tuning a strong dense retriever. we also find cocondenser yields performance close or similar to semi-supervised pretrained models that are several times larger. importantly, cocondenser provides a hands-off way to pre-train a very effective lm for dense retrieval. in particular, it effectively removes the effort for designing and testing pre-training as well as fine-tuning techniques. with our models, practitioners can use limited resources to train dense retrieval systems with state-of-the-art level performance. future works may also investigate integrating additional pre-training and/or fine-tuning methods to further improve performance.