[논문 리뷰] COCO-DR: Combating Distribution Shifts in Zero-Shot Dense Retrieval with Contrastive and Distributionally Robust Learning

Yue Yu, Chenyan Xiong외 3명|arXiv (Cornell University)|2022. 10. 27.

Domain Adaptation and Few-Shot Learning인용 수 11

한 줄 요약

COCO-DR는 대상 코퍼스에서의 연속 대비 미리 훈련을 통해 분포 이탈을 완화하고, 미세조정 중에 암묵적인 분포로 보수 최적화를 적용하는 제로샷 밀도 검색 방법이다. BEIR에서 최신 기술 수준의 성능을 달성하며, GPT-3 임베딩을 포함한 파arameter 수가 500배 이상 많은 모델들보다 뛰어나며, BERT Large 규모에서 파arameter 수의 0.17%만을 사용한다.

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 규모(110M 파arameter)에서 COCO-DR는 약 50배 더 많은 파arameter를 사용하는 GTR-XXL 및 CPT-L보다 BEIR 벤치마크에서 뛰어난 성능을 보였다.
• BERT Large 규모에서 COCO-DR는 파arameter 수가 175B에 이르는 CPT-XL을 뛰어넘었으며, 그 파arameter 수의 0.17%만을 사용하였다.
• 연속 대비 미리 훈련은 BioASQ에서 시퀀스 대비 손실을 50% 감소시켰고, 이는 제로샷 nDCG@10에서 22% 향상과 상관관계를 보였다.
• iDRO는 BEIR 작업에 가장 가까운 18개의 소스 질의 군집 중 12개에서 훈련 손실을 개선했으며, 소스 군집 성과 향상과 목표 작업 성능 향상 사이에 명확한 정적 상관관계를 보였다.
• COCO-DR는 표현의 균형 잡힌 정렬성과 균일성을 달성했으며, 이는 아블레이션 연구를 통해 더 나은 일반화와 관련이 있음을 입증했다.
• 모델은 BEIR의 18개 작업 전반에서 강력한 제로샷 일반화 성능을 보였으며, 작업별 적응이나 대규모 파arameter 확장을 통한 스케일업 없이도 강건성을 입증했다.