NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models

要約

デコーダ専用の大規模言語モデル (LLM) ベースの埋め込みモデルは、高密度ベクトルベースの検索などの汎用テキスト埋め込みタスクにおいて、BERT または T5 ベースの埋め込みモデルよりも優れたパフォーマンスを発揮し始めています。
この作業では、シンプルさと再現性を維持しながら、汎用性の高い埋め込みモデルとして LLM のパフォーマンスを大幅に向上させるための、さまざまなアーキテクチャ設計とトレーニング手順を備えた NV-Embed モデルを紹介します。
モデル アーキテクチャについては、プールされたエンベディングを取得するための潜在的なアテンション層を提案します。これにより、平均プーリングや LLM からの最後の トークン エンベディングを使用する場合と比較して、取得とダウンストリーム タスクの精度が一貫して向上します。
表現学習を強化するために、対照的なトレーニング中に LLM の因果的注意マスクを削除します。
モデルのトレーニングには、2 段階の対照的な命令調整方法を導入します。
まず、バッチ内のネガと厳選されたハード ネガの例を利用して、取得データセットに関する指示を含む対照的なトレーニングを適用します。
ステージ 2 では、さまざまな非検索データセットを命令チューニングにブレンドします。これにより、非検索タスクの精度が向上するだけでなく、検索パフォーマンスも向上します。
これらの技術を組み合わせた当社の NV-Embed モデルは、公開されているデータのみを使用して 69.32 という過去最高スコアを達成し、Massive Text Embedding Benchmark (MTEB) で第 1 位にランクされました (2024 年 5 月 24 日現在)。
タスクには、検索、再ランキング、分類、クラスタリング、および意味論的なテキスト類似タスクが含まれます。
特に、私たちのモデルは、MTEB ベンチマーク (BEIR とも呼ばれます) の 15 の検索タスクで最高スコア 59.36 も達成しています。
このモデルは https://huggingface.co/nvidia/NV-Embed-v1 でオープンソース化されます。

要約(オリジナル)

Decoder-only large language model (LLM)-based embedding models are beginning to outperform BERT or T5-based embedding models in general-purpose text embedding tasks, including dense vector-based retrieval. In this work, we introduce the NV-Embed model with a variety of architectural designs and training procedures to significantly enhance the performance of LLM as a versatile embedding model, while maintaining its simplicity and reproducibility. For model architecture, we propose a latent attention layer to obtain pooled embeddings, which consistently improves retrieval and downstream task accuracy compared to mean pooling or using the last token embedding from LLMs. To enhance representation learning, we remove the causal attention mask of LLMs during contrastive training. For model training, we introduce a two-stage contrastive instruction-tuning method. It first applies contrastive training with instructions on retrieval datasets, utilizing in-batch negatives and curated hard negative examples. At stage-2, it blends various non-retrieval datasets into instruction tuning, which not only enhances non-retrieval task accuracy but also improves retrieval performance. Combining these techniques, our NV-Embed model, using only publicly available data, has achieved a record-high score of 69.32, ranking No. 1 on the Massive Text Embedding Benchmark (MTEB) (as of May 24, 2024), with 56 tasks, encompassing retrieval, reranking, classification, clustering, and semantic textual similarity tasks. Notably, our model also attains the highest score of 59.36 on 15 retrieval tasks in the MTEB benchmark (also known as BEIR). We will open-source the model at: https://huggingface.co/nvidia/NV-Embed-v1.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google