[論文レビュー] 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 は、ターゲットコーパスにおける連続的対照的事前学習とファインチューニング時の暗黙的分布整合的最適化を用いて、分布シフトを是正するゼロショット密度検索手法であり、BEIR ベンチマークで最先端の性能を達成しており、GPT-3埋め込みを含むパラメータ数が500倍以上多いモデルを上回り、BERT Largeスケールでもパラメータのたった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 は、同じドキュメントからの2つのテキストシーケンスをポジティブペア、異なるドキュメントからのシーケンスをネガティブペアとして扱い、ターゲットコーパス上で連続的対照的事前学習（COCO）を実行する。
• これにより、シーケンス表現の整合性と一様性が向上し、ファインチューニングの前段階でモデルがターゲット分布に適応する。
• クエリの分布シフトに対処するため、COCO-DR はファインチューニング中に、埋め込み類似度に基づいてソースクエリをクラスタリングすることで、暗黙的分布整合的最適化（iDRO）を適用する。
• iDRO は勾配類似度を用いて各クラスタごとに損失を動的に再重み付けし、代表されないクエリクラスタにおけるモデルのロバスト性を向上させる。
• 本手法は、ターゲットクエリの明示的ラベルを必要とせず、分布適応とクエリクラスタのロバスト性の両方を同時に最適化する。
• モデルはソースタスクのラベル（例：MS MARCO）でのみファインチューニングされるが、BEIR の未学習ターゲットタスクに対しても効果的に一般化する。

実験結果

リサーチクエスチョン

• RQ1 クエリの意図や文書言語における分布シフトが、ゼロショット密度検索性能にどのように影響するか？
• RQ2 ターゲットコーパス上で連続的対照的事前学習を実施することで、ゼロショット検索におけるモデルの一般化性能が向上するか？
• RQ3 暗黙的分布整合的最適化（iDRO）は、稀なまたは未知のクエリクラスタにおいてモデルのロバスト性を向上させるか？
• RQ4 分布シフトに対処することで、小さなモデルがゼロショット検索において大規模モデルをどれほど上回れるか？

主な発見

• BERT Base スケール（110M パラメータ）で、COCO-DR はパラメータ数が約50倍多い GTR-XXL や CPT-L を BEIR ベンチマークで上回った。
• BERT Large スケールで、COCO-DR はパラメータ数が175B（CPT-XL）のモデルを、使用パラメータのたった0.17%で上回った。
• 連続的対照的事前学習により、BioASQ ではシーケンス対照的損失が50%低下し、これに相関してゼロショット nDCG@10 が22%向上した。
• iDRO は、BEIR タスクに最も近い18個のソースクエリクラスタのうち12個で訓練損失を改善し、ソースクラスタの向上とターゲットタスクの改善の間に明確な正の相関が見られた。
• COCO-DR は、表現のバランスの取れた整合性と一様性を達成しており、アブレーションスタディで示されるように、より良い一般化性能に寄与している。
• モデルは18のBEIRタスクにわたって強力なゼロショット一般化性能を示し、タスク固有の適応や大規模なパラメータスケーリングなしで、ロバスト性を示した。