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

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 1462 - 1479 December 7-11, 2022 ©2022 Association for Computational Linguistics COCO-DR: Combating Distribution Shifts in Zero-Shot Dense Retrieval with Contrastive and Distributionally Robust Learning Yue Yu1∗ Chenyan Xiong2 Si Sun3 Chao Zhang1 Arnold Overwijk2 1 Georgia Institute of Technology 2 Microsoft 3 Tsinghua University {yueyu, chaozhang}@gatech.edu, s-sun17@mails.tsinghua.edu.cn {chenyan.xiong, arnold.overwijk}@microsoft.com Abstract We present a new zero-shot dense retrieval (Ze- roDR) method, COCO-DR, to improve the gen- eralization ability of dense retrieval by com- bating 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 sam- ples from different source query clusters for im- proving model robustness over rare queries dur- ing fine-tuning. COCO-DR achieves superior average performance on BEIR, the zero-shot re- trieval benchmark. At BERTBase scale, COCO- DRBase outperforms other ZeroDR models with 60× larger size. At BERTLarge scale, COCO- DRLarge outperforms the giant GPT-3 embed- ding model which has 500× more parame- ters. Our analysis show the correlation of COCO-DR’s effectiveness in combating dis- tribution shifts and improving zero-shot accu- racy. Our code and model can be found at https://github.com/OpenMatch/COCO-DR. 1 Introduction Learning to represent and match queries and documents by embeddings, dense retrieval (DR) achieves strong performances in scenarios with sufficient training signals (Bajaj et al., 2016; Kwiatkowski et al., 2019). However, in many real world scenarios, obtaining relevance labels can be challenging due to the reliance on domain exper- tise, or even infeasible because of the strict privacy constraints. Deploying dense retrieval in these sce- narios becomes zero-shot (ZeroDR, Thakur et al. (2021)), which requires first training DR models on source tasks and then generalizing to target tasks with zero in-domain supervision (Izacard et al., 2022; Ni et al., 2021; Neelakantan et al., 2022). ∗Work partly done during Yue’s internship at Microsoft. 0.1 0.3 1.2 4.8 175.0 Parameters (in billions) 0.46 0.48 0.50 0.52 0.54 Avg. BEIR nDCG@10 BM25 DocT5query ColBERT Contriever GTR CPT-text COCO-DR Figure 1: The average nDCG@10 of COCO-DR versus large scale models on the 11 BEIR tasks selected in Neelakantan et al. (2022). X-axis is in log scale. ZeroDR poses great challenges to the general- ization ability of DR models under the distribution shift between source and target data (Gulrajani and Lopez-Paz, 2021; Wiles et al., 2022), 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 (Xin et al., 2022). In this work, we first analyze the distribution shifts in zero-shot dense retrieval. We illustrate the significant distribution shifts in both query in- tent and document language from the source to target tasks. After that, we show the strong correla- …

Content truncated to 500 words