[论文解读] COCO-DR: Combating Distribution Shifts in Zero-Shot Dense Retrieval with Contrastive and Distributionally Robust Learning

Yue Yu, Chenyan Xiong+3 more|arXiv (Cornell University)|Oct 27, 2022

Domain Adaptation and Few-Shot Learning被引用 11

一句话总结

COCO-DR 是一种零样本密集检索方法，通过在目标语料上进行连续对比预训练，并在微调过程中采用隐式分布鲁棒优化（iDRO），有效缓解了分布偏移问题。该方法在 BEIR 上实现了最先进性能，仅使用 GPT-3 嵌入模型 0.17% 的参数量，便在 BERT Large 规模下超越了参数量多出 500 倍的模型。

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 show the correlation between COCO-DR's effectiveness in combating distribution shifts and improving zero-shot accuracy. Our code and model can be found at \url{https://github.com/OpenMatch/COCO-DR}.

研究动机与目标

• 解决源训练任务与目标零样本检索场景之间的分布偏移问题。
• 在无领域内监督的低资源、分布外设置下，提升密集检索模型的泛化能力。
• 开发一种轻量化、统一的模型，确保在多样化检索任务中保持强大的零样本性能。
• 通过实证验证连续对比学习与隐式 DRO 在缓解分布偏移方面的有效性。

提出的方法

• COCO-DR 通过对目标语料进行连续对比预训练（COCO），将同一文档中的两个文本序列视为正样本对，不同文档中的序列视为负样本对。
• 该方法提升了序列表示的对齐性与均匀性，使模型在微调前即适应目标分布。
• 针对查询分布偏移，COCO-DR 在微调过程中应用隐式分布鲁棒优化（iDRO），通过嵌入相似性对源查询进行聚类。
• iDRO 利用梯度相似性动态重加权各聚类的损失，提升对低频查询聚类的鲁棒性。
• 该方法在无需目标查询显式标签的情况下，联合优化分布适应与查询聚类鲁棒性。
• 模型仅使用源任务标签（如 MS MARCO）进行微调，却能有效泛化至 BEIR 中未见的目标任务。

实验结果

研究问题

• RQ1 查询意图与文档语言的分布偏移如何影响零样本密集检索的性能？
• RQ2 在目标语料上进行连续对比预训练是否能提升零样本检索中的模型泛化能力？
• RQ3 隐式分布鲁棒优化是否能提升模型在罕见或未见查询聚类上的鲁棒性？
• RQ4 通过缓解分布偏移，小型模型在多大程度上可超越大规模模型的零样本检索性能？

主要发现

• 在 BERT Base 规模（1.1 亿参数）下，COCO-DR 在 BEIR 基准上超越了使用约 50 倍参数的 GTR-XXL 和 CPT-L 模型。
• 在 BERT Large 规模下，COCO-DR 在部分 BEIR 任务上超越了参数量达 1750 亿的 CPT-XL 模型，且仅使用其 0.17% 的参数量。
• 连续对比预训练使 BioASQ 上的序列对比损失降低 50%，与零样本 nDCG@10 提升 22% 呈显著正相关。
• iDRO 在 18 个最接近 BEIR 任务的源查询聚类中，有 12 个实现了训练损失降低，且源聚类性能增益与目标任务性能提升之间存在明确正相关。
• COCO-DR 实现了表示学习中对齐性与均匀性的良好平衡，该特性与更优泛化能力相关，经消融实验证实。
• 该模型在 18 个 BEIR 任务上均展现出强大的零样本泛化能力，证明其无需任务特定微调或大规模参数扩展即可保持鲁棒性。