COCO - DR : Combating the Distribution Shift in Zero-Shot Dense Retrieval with Contrastive and Distributionally Robust Learning (2022)

Abstract: We present a new zero-shot dense retrieval (ZeroDR) method, COCO-DR, to improve the generalization ability of dense retrieval by combating the distribution shifts between source training tasks and target scenarios. To mitigate the impact of document differences, COCO-DR continues pretraining the language model on the target corpora to adapt the model to target distributions via COtinuous COtrastive learning. To prepare for unseen target queries, COCO-DR leverages implicit Distributionally Robust Optimization (iDRO) to reweight samples from different source query clusters for improving model robustness over rare queries during fine-tuning. COCO-DR achieves superior average performance on BEIR, the zero-shot retrieval benchmark. At BERT_Base scale, COCO-DR Base outperforms other ZeroDR models with 60x larger size. At BERT_Large scale, COCO-DR Large outperforms the giant GPT-3 embedding model which has 500x more parameters. Our analysis shows the correlation between COCO-DR’s effectiveness in combating distribution shifts and improving zero-shot accuracy. Our code and model can be found at https://github.com/OpenMatch/COCO-DR .

## introduction learning to represent and match queries and documents by embeddings, dense retrieval (dr) achieves strong performances in scenarios with sufficient training signals @xcite . however, in many real world scenarios, obtaining relevance labels can be challenging due to the reliance on domain expertise, or even infeasible because of the strict privacy constraints. deploying dense retrieval in these scenarios becomes zero-shot (zerodr, @xcite ), which requires first training dr models on source tasks and then generalizing to target tasks with zero in-domain supervision @xcite @xcite . zerodr poses great challenges to the generalization ability of dr models under the distribution shift between source and target data @xcite , as it requires the alignment between queries and their relevant documents in the embedding space. it is much harder to generalize than standard classification or ranking tasks, where a robust decision boundary is sufficient @xcite . ## related work earlier research has explored various ways to learn representations for retrieval @xcite . recently, with pretrained language models @xcite , hard training negative selection @xcite , and retrieval-oriented pretraining @xcite , dense retrieval has shown strong advantages over sparse retrieval methods, although the advantages are more observed in supervised settings than zeroshot scenarios @xcite . one research direction to improve zero-shot dense retrieval is bringing in domain adaption techniques. @xcite employ domain invariant learning to narrow the representation gap between source and target domains. @xcite and @xcite generate pseudo labels for each target task to train in-domain dr models. these techniques employ one specially trained retrieval model for each target task and improve zero-shot retrieval accuracy. another way to improve zerodr is to scale up model size and source training data. @xcite and @xcite leverage models with billions of parameters (t5-xxl and gpt-3) and large-scale training data to increase the generalization capacity of dr model. @xcite and @xcite enlarge the size of training data with retrieval-oriented pretraining tasks. as illustrated in figure 1 , the benefit of scale follows the scaling law of language models @xcite : a linear increment of zero-shot accuracy requires exponentially more training data and model parameters. combining dense models with sparse retrieval yields better zero-shot retrieval performances on beir @xcite . the reranking models, using stronger cross-encoders, can be used as teachers to improve the robustness of dense retrieval models @xcite . more generally speaking, continuous pretraining and distributionally robust optimization (dro) are two techniques for improving model generalization on other applications. continuous pretraining bert's masked language modeling tasks on target domain corpora have shown benefits on both language tasks @xcite and the reranking step of search systems @xcite . the benefits of dro are more ambivalent @xcite and are more observed when explicit group partitions are available @xcite @xcite . ## distribution shifts in dense retrieval in this section, we first introduce the preliminaries of dense retrieval. then we discuss the standard zero-shot dense retrieval settings and study the impact of distribution shifts on zerodr accuracy. ## coco-dr method to combat the distribution shifts from training source to zero-shot targets, coco-dr introduces two training techniques: continuous contrastive pretraining (coco) and implicit distributionally robust optimization (idro). the first continuously pretrains the language model on target corpora to handle document distribution shifts. the latter improves the model robustness during fine-tuning, which then lead to better generalization for unseen target queries. this section describes these two components in detail. ## experiments in this section, we first describe our experiment setups and evaluate coco-dr. then we analyze the efficacy of coco and idro. ## conclusion coco-dr improves zerodr accuracy by combating the distribution shifts using continuous con-trastive learning and implicit distributionally robust optimization. coco helps models better capture the sequence representations of target corpora in pretraining. implicit dro improves model robustness by reweighting query clusters in fine-tuning. coco-dr achieves strong zero-shot performance while maintaining a lightweight system with one unified model for all 18 target tasks. different than prior works that scaling up the dr model to billions of parameters (e.g. cpt-text), we provide a more efficient and sustainable way to improve the zero-shot generalization ability. our analyses observed clear correlations on coco-dr's ability to mitigate the distribution shifts and to generalize. better zerodr accuracy is observed on tasks where continuous contrastive learning has a lower pretraining loss, and where idro successfully identifies and improves source query clusters similar to target queries.