使用LLM2Vec将Llama 3转变为一个嵌入模型

文章目录

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

Github：https://github.com/McGill-NLP/llm2vec

笔记：https://kaitchup.substack.com/p/notebooks

由DALL-E生成

嵌入模型是大型语言模型（LLMs）中检索增强生成（RAG）的关键组成部分。它们对用户编写的知识库和查询进行编码。

使用与LLM相同领域训练或微调的嵌入模型可以显著改善RAG系统。然而，找到或训练这样的嵌入模型通常是一项困难的任务，因为领域内的数据通常很少。

在本文中，我将展示如何使用LLM2Vec将LLM转变为文本嵌入模型。我们将看到如何使用Llama 3来创建一个RAG系统，该系统不需要除Llama 3之外的任何其他模型。

LLM2Vec：您的LLM也是一个嵌入模型

LLM2Vec 在这篇论文中提出：

LLM2Vec: 大型语言模型实际上是强大的文本编码器 https://arxiv.org/pdf/2404.05961.pdf

官方实现可在GitHub上找到：

LLM2Vec是如何工作的？

LLMs是通过因果语言建模损失进行训练的。它们被训练以预测给定一系列标记的下一个标记。由于训练示例是整个标记序列，因此在序列上应用了因果注意力掩码，以便当模型学会预测一个标记时，序列中的所有后续标记都被屏蔽，不影响注意力计算。例如，如果我们有以下序列：

The cat is sleeping in the kitchen <eos>

在训练期间，将依次应用以下注意力掩码（在掩码下，我仅显示未被屏蔽的标记）：

1 0 0 0 0 0 0 0 0
The 

1 1 0 0 0 0 0 0
The cat

1 1 1 0 0 0 0 0
The cat is 

1 1 1 1 0 0 0 0
The cat is sleeping

1 1 1 1 1 0 0 0

The cat is sleeping in 

1 1 1 1 1 1 0 0
The cat is sleeping in the 

1 1 1 1 1 1 1 0
The cat is sleeping in the kitchen

1 1 1 1 1 1 1 1
The cat is sleeping in the kitchen <eos>

零表示该标记不会被考虑进行注意力计算。

另一方面，文本嵌入模型被训练以编码整个序列标记，并且是双向的，即它们从左到右和从右到左编码序列。

LLM2Vec方法中将LLM转换为嵌入模型的初始步骤涉及用全为1的矩阵替换仅解码器LLM中使用的因果注意力掩码。这种改变使得序列中的每个标记都可以与所有其他标记互动，有效地将其转变为双向LLM。然而，仅解码器LLM未经训练以编码未来标记，因此，这种简单修改会降低表示的质量。尽管如此，可以训练模型有效地利用其新的双向注意力能力。

他们建议采用新的掩码下一个标记预测（MNTP）目标。MNTP将下一个标记预测目标与BERT使用的掩码语言建模合并。为了实现这一点，我们取一个任意序列x = (x1, x2, . . . , xN)，掩盖一部分输入标记，然后训练模型使用过去和未来上下文来预测这些被屏蔽的标记。重要的是，在预测在位置i被屏蔽的标记时，损失是使用来自前一位置i−1的标记表示的logits计算的，而不是来自被屏蔽位置本身，就像对BERT模型所做的那样。

将LLM转换为双向模型，然后进行MNTP训练可以将任何仅解码器LLM调整为编码器。然而，这些步骤可能无法充分解决序列表示的需求。事实上，像BERT这样的双向编码器通常还在它们的预训练中包含一个下一个句子预测任务，而LLMs并未接受此任务的训练。LLM2Vec通过将每个输入序列在模型中处理两次，每次使用不同的随机选择的dropout掩码，产生相同序列的两个不同表示来解决这一缺失能力。训练目标是提高这两个表示之间的相似性，同时减少它们与同一训练批次中不同序列的表示之间的相似性。他们将这最后一步称为“无监督对比学习”（SimCSE）。

LLM2Vec背后的完整过程由论文中的这个图示说明：

来源（CC-BY）

他们评估了LLM2Vec生成的模型在各种任务中的表现，并显示它们可以胜过标准文本嵌入模型。您可以在论文的第3和第4节中找到结果。

使用LLM2Vec将Llama 3转变为文本嵌入模型

我的笔记本展示了如何将 Llama 3 转换为嵌入模型，点击此处查看：

获取笔记本 (#65) https://kaitchup.substack.com/p/notebooks

使用 LLM2Vec 将 LLM 转换为文本嵌入模型相当简单。

首先，安装以下软件包：

pip install llm2vec
pip install flash-attn --no-build-isolation

llm2vec 软件包将 LLM 转换为嵌入模型。flash-attn 是 FlashAttention 的软件包。虽然不是必需的，但可以加快使用 MNTP 训练的速度。它仅适用于安培（Ampere）世代的最新 GPU（RTX 3xxx/4xxx、A100、H100 等）。

然后，转换本身由以下代码执行：

import torch
from llm2vec import LLM2Vec

l2v = LLM2Vec.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    device_map="cuda" if torch.cuda.is_available() else "cpu",
    torch_dtype=torch.bfloat16,
)

l2v.save("Llama-3-8B-Emb")

“torch_dtype=torch.bfloat16” 是为了能够在 24 GB 的 GPU 上运行转换。如果不设置它，模型将比具有 float32 参数的原始模型更大。

我将此模型推送到了 hub，以防您不想自己进行转换：

此模型已准备就绪。您可以将其插入 RAG pipeline 中。但在大多数情况下，它的性能不如标准的嵌入模型。

我们需要用 MNTP 目标训练 Llama 3 8B。作者还提供了一个脚本来执行此操作：

目前它支持具有 Llama 和 Mistral 架构的模型。

要使用它，请在本地克隆存储库：

git clone https://github.com/McGill-NLP/llm2vec.git

该脚本期望一个参数，即 JSON 格式的配置文件。他们在这里提供了几个示例：

对于 Llama 3 8B，我创建了以下配置：

JSON_CONFIG='''
{
    "model_name_or_path": "meta-llama/Meta-Llama-3-8B",
    "dataset_name": "wikitext",
    "dataset_config_name": "wikitext-103-raw-v1",
    "per_device_train_batch_size": 1,
    "per_device_eval_batch_size": 1,
    "gradient_accumulation_steps": 16,
    "do_train": true,
    "do_eval": true,
    "max_seq_length": 512,
    "mask_token_type": "blank",
    "data_collator_type": "all_mask",
    "mlm_probability": 0.8,
    "overwrite_output_dir": true,
    "output_dir": "Llama-3-8B-llm2vec-MNTP-Emb",
    "evaluation_strategy": "steps",
    "eval_steps": 100,
    "save_steps": 200,
    "stop_after_n_steps": 1000,
    "lora_r": 16,
    "gradient_checkpointing": true,
    "torch_dtype": "bfloat16",
    "attn_implementation": "flash_attention_2"
}
'''

with open("mntp_config.json", 'w') as f:
  f.write(JSON_CONFIG)

该脚本使用 bfloat16 参数加载模型。我将每个设备的批次大小设置为一，以便训练适合 24 GB GPU。

然后，我们可以开始 MNTP 训练：

python llm2vec/experiments/run_mntp.py mntp_config.json

使用 24 GB GPU（例如 Google Colab 的 L4）进行三次迭代的训练应该需要 4 天时间。如果您等不及那么长的时间，一个迭代可能就足够了。

经过 MNTP 训练，模型应该会产生更好的结果，特别是对于检索任务。

对于最后一步 SimCSE 训练，作者尚未发布他们的代码，但提到他们会发布。

为 RAG 设置 Llama 3 文本嵌入模型

在前一节创建的嵌入模型可以在 RAG pipeline 中使用。例如，您可以使用 sentence-transformers（Apache 2.0 许可证）加载它。

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("kaitchup/Llama-3-8B-llm2vec-Emb")

如果您使用 LlamaIndex（MIT 许可证），可以设置 HuggingFaceEmbedding 模型：

Settings.embed_model = HuggingFaceEmbedding(model_name="kaitchup/Llama-3-8B-llm2vec-Emb", device='cpu')

我将 device=‘cpu’，但使用 CPU 会使 RAG 系统变慢。您可以删除此参数以在 GPU 上运行它。但请注意，模型是以完整精度加载的。它不适合消费级 GPU。

结论

使用 LLM2Vec，我们现在可以将 LLM 用作文本嵌入模型。转换简单而快速。在 RAG 系统中同时用一个模型进行生成和检索非常吸引人，因为我们不需要再搜索额外的嵌入模型。