NV-Embed：汎用埋め込みモデルとしての大規模言語モデルを訓練するための改良技術

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

May 27, 2024

著者: Chankyu Lee, Rajarshi Roy, Mengyao Xu, Jonathan Raiman, Mohammad Shoeybi, Bryan Catanzaro, Wei Ping

cs.AI

要旨

Decoder-onlyの大規模言語モデル（LLM）ベースの埋め込みモデルは、BERTやT5ベースの埋め込みモデルを、密なベクトルベースの検索を含む汎用テキスト埋め込みタスクにおいて上回り始めています。本研究では、LLMを汎用埋め込みモデルとしての性能を大幅に向上させながら、そのシンプルさと再現性を維持するために、様々なアーキテクチャ設計とトレーニング手順を備えたNV-Embedモデルを紹介します。モデルアーキテクチャに関しては、プールされた埋め込みを取得するための潜在注意層を提案し、これにより平均プーリングやLLMの最後の<EOS>トークン埋め込みを使用する場合と比較して、検索および下流タスクの精度が一貫して向上します。表現学習を強化するために、LLMの因果的注意マスクを対照学習中に除去します。モデルトレーニングに関しては、2段階の対照的指示チューニング手法を導入します。まず、検索データセットに対して指示を用いた対照学習を適用し、バッチ内のネガティブ例と精選されたハードネガティブ例を活用します。第2段階では、様々な非検索データセットを指示チューニングに組み込み、非検索タスクの精度を向上させるだけでなく、検索性能も改善します。これらの技術を組み合わせることで、公開されているデータのみを使用した我々のNV-Embedモデルは、Massive Text Embedding Benchmark（MTEB）（2024年5月24日現在）において、検索、再ランキング、分類、クラスタリング、意味的テキスト類似性タスクを含む56のタスクで、69.32という記録的なスコアを達成し、第1位にランクされました。特に、我々のモデルはMTEBベンチマーク（BEIRとも呼ばれる）の15の検索タスクにおいても59.36という最高スコアを獲得しました。モデルは以下のURLでオープンソース化されます：https://huggingface.co/nvidia/NV-Embed-v1。

English

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 <EOS> 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.