[論文レビュー] Towards General Text Embeddings with Multi-stage Contrastive Learning

Zehan Li, Xin Zhang+4 more|arXiv (Cornell University)|Aug 7, 2023

Topic Modeling被引用数 57

ひとこと要約

GTEは多様なオープンデータの混合から多段階対照学習を用いて汎用的なテキスト埋め込みモデルを訓練し、比較的小さなモデルサイズでテキスト、コード、検索タスクの零-shotおよび監督付き性能を強力に達成します。

ABSTRACT

We present GTE, a general-purpose text embedding model trained with multi-stage contrastive learning. In line with recent advancements in unifying various NLP tasks into a single format, we train a unified text embedding model by employing contrastive learning over a diverse mixture of datasets from multiple sources. By significantly increasing the number of training data during both unsupervised pre-training and supervised fine-tuning stages, we achieve substantial performance gains over existing embedding models. Notably, even with a relatively modest parameter count of 110M, GTE$_ ext{base}$ outperforms the black-box embedding API provided by OpenAI and even surpasses 10x larger text embedding models on the massive text embedding benchmark. Furthermore, without additional fine-tuning on each programming language individually, our model outperforms previous best code retrievers of similar size by treating code as text. In summary, our model achieves impressive results by effectively harnessing multi-stage contrastive learning, offering a powerful and efficient text embedding model with broad applicability across various NLP and code-related tasks.

研究の動機と目的

• 多様なデータソースを横断して、多段階対照学習で訓練された統一テキスト埋め込みモデルを開発する。
• 一般化を向上させるために、大規模な自己教師あり事前学習を活用し、続いて監督付きファインチューニングを行う。
• テキストおよびコードタスクで、より大規模なモデルやタスク特化型埋め込みモデルと競争力がある、あるいは優位性を示す。

提案手法

• テキスト埋め込みを得るために、平均プーリングを用いたデュアルエンコーダー Transformer バックボーンを使用する。
• 多様なオープンソースからの約8億のテキストペアに対して、自己教師あり対照学習で事前訓練する。
• 複数タスクからの約300万の注釈付きテキストトリプルを用いた監督付き対照学習でファインチューニングする。ハードネガティブを含む。
• バッチ内のクエリ/ドキュメントおよび双方向用語を拡張してネガティブを増やし、対照ロスを改善する。
• ソースのバランスをとるために、複数のソースから多項分布でデータバッチをサンプリングする（α = 0.5）。
• タスク固有のプロンプトを用いずに、零-shot検索、MTEB、コード検索ベンチマークを横断して評価する。

実験結果

リサーチクエスチョン

• RQ1 多様なオープンデータを横断して多段階対照学習で学習された一般的なテキスト埋め込みモデルは、標準ベンチマークでタスク専用またはより大きな埋め込みモデルを上回ることができるか？
• RQ2 データの多様性、バッチサイズ、およびモデル規模は、自己教師あり事前学習と監督付きファインチューニングの性能にどう影響しますか？
• RQ3 言語特化のファインチューニングを行わずにテキスト埋め込みとして訓練した場合、コード検索は効果的に改善されますか？

主な発見

• GTEbase (110M) は OpenAI の embedding API を上回り、巨大なテキスト埋め込みベンチマークで 10 倍以上大きいモデルの中には勝るものもある。
• オープンデータでの自己教師あり事前学習と監督付きファインチューニングは、MTEB タスク、BEIR、コード検索の分野で強力な性能を発揮する。
• GTEbase は、タスク特定のプロンプトを用いずに零-shotのテキスト検索やテキスト分類タスクで、大規模モデルと同等またはそれ以上の性能を示す。
• GTEbase は、同等サイズの最新コード検索システムを大きく上回り、CodeSearchNet で言語特化のファインチューニングなしで成果を出す。
• データ多様性とモデルサイズを拡大すると線形に性能が向上し、ネガティブのバッチサイズの飽和は約1万件であり、多段階訓練は単段階訓練を上回る。