How to Train Your Dragon: Diverse Augmentation Towards Generalizable Dense Retrieval (2023)

Abstract: Various techniques have been developed in recent years to improve dense retrieval (DR), such as unsupervised contrastive learning and pseudo-query generation. Existing DRs, however, often suffer from effectiveness tradeoffs between supervised and zero-shot retrieval, which some argue was due to the limited model capacity. We contradict this hypothesis and show that a generalizable DR can be trained to achieve high accuracy in both supervised and zero-shot retrieval without increasing model size. In particular, we systematically examine the contrastive learning of DRs, under the framework of Data Augmentation (DA). Our study shows that common DA practices such as query augmentation with generative models and pseudo-relevance label creation using a cross-encoder, are often inefficient and sub-optimal. We hence propose a new DA approach with diverse queries and sources of supervision to progressively train a generalizable DR. As a result, DRAGON, our Dense Retriever trained with diverse AuGmentatiON, is the first BERT-base-sized DR to achieve state-of-the-art effectiveness in both supervised and zero-shot evaluations and even competes with models using more complex late interaction.

## introduction bi-encoder based neural retrievers allow documents to be pre-computed independently of queries and stored, enabling end-to-end retrieval among huge corpus for downstream knowledgeintensive tasks @xcite . recently, @xcite show that bi-encoder retrievers still underperform bm25 in real-world scenarios, where training data is scarce. one potential solution is to design more expressive representations to capture more fine-grained token-level information; e.g., splade++ @xcite and col-bertv2 @xcite in figure 1 . however, these designs add complexity and latency to retrieval systems @xcite . by contrast, dense retrieval (dr) is a simpler biencoder retriever that maps queries and documents into low-dimensional vectors and computes text similarity through a simple dot product. top-k retrieval can be performed directly using ann search libraries @xcite . recently, various methods have been proposed to improve dr effectiveness while keeping its simple architecture, such as pre-training @xcite , query augmentation @xcite , and distillation @xcite . for example, pre-training on ms marco corpus improves accuracy in the fully supervised setting while leveraging other corpora can improve transfer in the zero-shot setting. however, improvement in one setting is only achieved at the expense of the other. figure 1 plots existing state-of-the-art dr models with respect to their effectiveness on these two axes, which presents a clear tradeoff between supervised and zero-shot effectiveness (the blue line). the only exception, gtr-xxl @xcite , breaks the effectiveness tradeoff at the expense of efficiency (i.e., query encoding), which leverages very large pre-trained models with 40 times more parameters. this effectiveness tradeoff prompts some to hypothesize that we have fully exploited the capacity of bert-basesized dr model @xcite and explore how to cleverly increase model parameters without sacrificing retrieval efficiency. for example, the recent work @xcite proposes to train one expert dense retriever for each specific scenario, resulting in slow adaptation to real-world applications @xcite . in this work, we contradict this hypothesis and show that a generalizable dr can indeed be trained to achieve state-of-the-art effectiveness in both supervised and zero-shot evaluations without increasing model size. to this end, we first investigate the important factors contributing to the recent progress of dr. for example, dr seems to gain zero-shot transfer capability from pre-training on large-scale and diverse training queries @xcite while knowledge distillation can improve the supervision quality by automatically identifying relevant passages which are not labeled by humans @xcite . to better understand these approaches, we devise a unified framework of data augmentation (da) for contrastive learning. under the framework, the previous work can be viewed as da with different recipes of query augmentation and relevance label augmentation shown in guided by a detailed empirical exploration along the space of our da framework, we find the following: (1) for relevance label augmentation, we identify that the key to training a generalizable dense retriever is to create diverse relevance labels for each query, for which we use multiple retrievers instead of a strong cross encoder; (2) with such diverse relevance labels, dense retrievers can be trained effectively using cheap and large-scale augmented queries (e.g., cropped sentences from a corpus) instead of the more expensive neural generative queries. this finding opens the door to further building cheap but useful training data in scale for dr in the future. finally, we find that it is suboptimal for a dense retriever to learn the ## background in this section, we first introduce the retrieval task and contrastive learning for dense retrieval. we then provide a unified framework for understanding recent approaches to improve dense retrieval training as instances of data augmentation. ## pilot studies on data augmentation in this section, we first discuss the exploration space of data augmentation (da) based on the framework in section 2.2 and then conduct empirical studies on how to better train a dense retriever. based on the empirical studies, we propose our da recipe to train dragon, a dense retriever with diverse augmentation. ## conclusion we present dragon, a dense retriever trained with diverse augmentation and a unified framework of data augmentation (da) to understand the recent progress of training dense retrievers. based on the framework, we extensively study how to improve dense retrieval training through query and relevance label augmentation. our experiments uncover some insights into training a dense retriever, which contradicts common wisdom that cross encoder is the most effective teacher and human-like queries are the most suitable training data for dense retrieval. then, we propose a diverse data augmentation recipe, query augmentation with the mixture of sentence cropping and generative queries, and progressive relevance label augmentation with multiple teachers. with our proposed recipe of da, dragon is the first to demonstrate that a single bert-basesized dense retriever can achieve state-of-the-art effectiveness in both supervised and zero-shot retrieval tasks. we believe that dragon can serve as a strong foundation retrieval model for domain adaptation retrieval tasks @xcite or the existing retrieval augmented language models @xcite @xcite .