Towards General Text Embeddings with Multi-stage Contrastive Learning

要約

多段階の対照学習でトレーニングされた汎用テキスト埋め込みモデルである GTE を紹介します。
さまざまな NLP タスクを 1 つの形式に統合するという最近の進歩に合わせて、複数のソースからのデータセットの多様な混合に対する対照学習を採用することで、統一されたテキスト埋め込みモデルをトレーニングします。
教師なし事前トレーニング段階と教師あり微調整段階の両方でトレーニング データの数を大幅に増やすことにより、既存の埋め込みモデルと比較して大幅なパフォーマンス向上を実現します。
特に、パラメータ数が 1 億 1,000 万と比較的控えめであっても、GTE$_\text{base}$ は OpenAI が提供するブラックボックス埋め込み API を上回り、大規模テキスト埋め込みベンチマークでは 10 倍大きいテキスト埋め込みモデルさえも上回っています。
さらに、各プログラミング言語を個別に微調整する必要がなく、コードをテキストとして扱うことで、私たちのモデルは同様のサイズのこれまでの最高のコード取得ツールよりも優れたパフォーマンスを発揮します。
要約すると、私たちのモデルは、多段階の対照学習を効果的に活用することで素晴らしい結果を達成し、さまざまな NLP およびコード関連のタスクに幅広く適用できる強力で効率的なテキスト埋め込みモデルを提供します。

要約(オリジナル)

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$_\text{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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google