NV-Embed重塑文本嵌入艺术

论文：https://arxiv.org/pdf/2405.17428

代码：未开源

模型：https://huggingface.co/nvidia/NV-Embed-v1

NVIDIA团队提出NV-Embed，采用多种架构设计和训练流程，显著提高LLM作为多功能嵌入模型的性能，同时保持其简单性和可重现性。

1. 对于模型架构，提出了一个潜在的注意力层来获取汇总嵌入，与使用LLMs中的平均汇总或最后一个标记嵌入相比，它始终改善了检索和下游任务的准确性。

2. 为了增强表示学习，们在对比训练期间去除了LLMs的因果注意力掩码。

3. 对于模型训练，引入了一个两阶段的对比指导调整方法。首先，它在检索数据集上使用带有指导的对比训练，利用批内负样本和策划的困难负样本。在第二阶段，它将各种非检索数据集融入到指导调整中，不仅提高了非检索任务的准确性，还改善了检索性能。

通过结合这些技术，NV-Embed模型仅使用公开可用的数据，在 Massive Text Embedding Benchmark（MTEB）上取得了69.32的创纪录高分，截至2024年5月24日，排名第一，涵盖了56个任务，包括检索、重新排序、分类、聚类和语义文本相似性任务。值得注意的是，NV-Embed模型在MTEB基准测试中的15个检索任务中也取得了59.36的最高分（也称为BEIR）。

摘要（Abstract）

• 目的：提高大型语言模型（LLM）在文本嵌入任务上的性能。
• 方法：提出了NV-Embed模型，包含新架构和训练过程。
• 结果：使用潜在注意力层和两阶段对比指令调整方法，在多个基准测试上取得优异成绩。
• 开源：模型将在Hugging Face上开源。

1. 引言（Introduction）

• 文本嵌入的重要性：作为许多下游任务的基础，如检索、重排等。
• 现有模型的局限性：基于BERT或T5的模型虽然长期占主导地位，但存在局限性。
• NV-Embed的优势：不依赖专有数据，使用公开数据集训练，具有更好的泛化能力和性能。

2. 相关工作（Related Work）

• 2.1 双向嵌入模型：基于BERT或T5，通过对比学习进一步预训练。
• 2.2 解码器仅LLM基础的嵌入模型：先前工作显示解码器仅LLM在嵌入任务上可能表现不佳，但最新研究开始显示其潜力。

3. 方法（Method）

• 3.1 双向注意力：通过移除因果注意力掩码，允许模型在自回归文本生成中双向学习，增强表示能力。
• 3.2 潜在注意力层：提出一种新的池化方法，使用潜在数组和多头注意力机制，以获得更好的序列表示。
  • 查询（Q）和键值（KV）：使用解码器最后层的隐藏状态作为查询，潜在数组作为键值输入。
  • 输出（O）：通过softmax归一化的QKT矩阵与V相乘，并通过MLP层进一步处理。
• 3.3 两阶段指令调整：
  • 第一阶段：在检索数据集上使用对比训练，利用批内负样本和人工筛选的硬负样本。
  • 第二阶段：混合非检索数据集进行指令调整，不使用批内负样本，以提高分类、聚类和语义文本相似性任务的准确性，同时意外地增强了检索性能。

4. 训练数据（Training Data）

• 公共检索数据集：使用了多个数据集，如MS MARCO、HotpotQA等，并采用另一种基于编码器的嵌入模型来筛选硬负样本。
• 公共非检索数据集：包括分类、聚类和语义相似性任务的数据集，采用与检索数据集相同的格式进行对比训练。

5. 实验（Experiments）

• 5.1 实验细节：
  • 使用了低秩适应（LoRA）方法对NV-Embed模型进行高效微调。
  • 选择了Mistral 7B作为基础模型，并对其进行了端到端的对比损失训练。
• 5.2 MTEB结果：在56个任务上评估了NV-Embed模型，并与前沿模型进行了比较，显示出新记录的高分。
• 5.3 消融研究：
  • 因果注意力与双向注意力：证明了双向注意力在嵌入任务上更有效。
  • 池化方法：比较了不同的池化技术，发现潜在注意力层在检索和分类任务上特别有效。

  本文首发于公众号：AI智语狂想，欢迎关注。