NV-Embed：将 LLM 训练为通用嵌入模型的改进技术

基于仅解码器 LLM 的嵌入模型开始在通用文本嵌入任务（包括基于稠密向量的检索）中超越基于 BERT 或 T5 的嵌入模型。在本工作中，我们提出了 NV-Embed，结合架构设计、训练流程和精选数据集，在保持简单性和可复现性的同时，显著增强了 LLM 作为通用嵌入模型的性能。在模型架构方面，我们提出了一个潜在注意力层来获取池化嵌入，与均值池化或使用 LLM 的最后一个 <EOS> 词嵌入相比，它持续提升了检索和下游任务的准确率。为了增强表示学习，我们在对比训练期间移除了 LLM 的因果注意力掩码。在训练算法方面，我们引入了一种两阶段对比指令微调方法。它首先在检索数据集上应用带有指令的对比训练，利用批内负样本和精选的困难负样本。在第二阶段，它将各种非检索任务混合到指令微调中，这不仅提高了非检索任务的准确率，也改善了检索性能。在训练数据方面，我们利用困难负样本挖掘、合成数据生成和现有公开数据集来提升嵌入模型的性能。通过结合这些技术，我们的 NV-Embed-v1 和 NV-Embed-v2 模型在 56 项任务上取得了 MTEB 排行榜的第一名（分别截至 2024 年 5 月 24 日和 8 月 30 日），证明了所提方法随时间的持续有效性。它还在 AIR 基准测试的长文档部分取得了最高分，在问答部分取得了第二高分，该基准测试涵盖了 MTEB 之外的一系列域外信息检索主题。我们进一步提供了通用嵌入模型的模型压缩技术分析。