LLM2Vec介绍和将Llama 3转换为嵌入模型代码示例

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

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

但是这篇论文LLM2Vec，可以将任何的LLM转换为文本嵌入模型，这样我们就可以直接使用现有的大语言模型的信息进行RAG了。

嵌入模型和生成模型

嵌入模型主要用于将文本数据转换为数值形式的向量表示，这些向量能够捕捉单词、短语或整个文档的语义信息。这些向量表示也被称为嵌入（embeddings），可以用于各种下游任务，如文本分类、搜索、相似度计算等。

最有名的嵌入模型就是BERT是一个典型的encoder-only模型

生成模型则设计用来基于训练数据生成新的数据实例。在NLP中，这通常意味着生成文本。这类模型能够学习到数据的分布，并能创造出符合这一分布的新实例，如新的句子或文档。

如GPT系列，通常是decoder-only模型。

这两种架构在设计和应用上有所不同：

1. BERT (Encoder-only)：BERT利用双向Transformer编码器，这意味着它在处理文本时可以同时考虑前面和后面的上下文。这种双向上下文理解使得BERT非常适合用于各种理解任务，如问答、自然语言推理和实体识别等。
2. LLM (Decoder-only)：如GPT系列模型，通常采用单向Transformer解码器。这意味着在生成文本时，每个新词只能基于前面的词生成。这种结构适合于文本生成任务，如文本续写、自动编写程序代码等。

在论文中对encoder-only和decoder-only模型的特点进行了讨论，特别是在解释为什么将decoder-only的大型语言模型（LLM）转换为有效的文本编码器时。论文指出了几个关键点：

1. Decoder-only模型的局限性：这些模型，如GPT系列，使用因果（单向）注意力机制。这意味着在生成文本时，每个标记只能看到它之前的标记。这种结构虽然适合文本生成，但在需要丰富的上下文信息（如文本嵌入任务）时可能不够理想。
2. 克服Decoder-only模型的限制：论文中提出的LLM2Vec方法特别通过几个步骤来克服这些限制，包括启用双向注意力，这使得模型能够在处理文本时同时考虑前后文本，从而生成更丰富的上下文表示。
3. 与Encoder-only模型的对比：论文比较了使用LLM2Vec转换的decoder-only模型与传统的encoder-only模型（如BERT）。结果显示，通过适当的转换和训练，原本设计为decoder-only的模型在多个文本嵌入任务上能够达到甚至超过传统encoder-only模型的性能。

这表明，尽管decoder-only和encoder-only模型在设计和功能上有本质的不同，但通过创新的方法可以扩展或改变这些模型的能力，使它们适应更广泛的应用场景。

其实我们可以将这篇论文的重点简单的理解为，如何将一个decoder-only的模型快速并且无损的转换成一个encoder-only模型。

LLM2Vec

在论文中提出了一种名为LLM2Vec的方法，用于将仅解码器的大型语言模型（LLM）转换为强大的文本编码器。这种方法包括三个简单的步骤：1）启用双向注意力；2）蒙版下一个标记预测；3）无监督对比学习。这种转换不需要标记数据，且在数据和参数上都非常高效。

具体来说，研究中首先解决了LLM在文本嵌入任务中由于其因果关注机制而受限的问题，该机制仅允许标记与其前面的标记交互。通过启用双向注意力，每个标记能够访问序列中的所有其他标记，从而转换为双向LLM。然后，通过蒙版下一个标记预测（MNTP），调整模型以利用其双向注意力。最后，应用无监督对比学习以改进序列表示。

这些步骤的组合不仅在单词级任务上提升了模型的性能，还在大规模文本嵌入基准（MTEB）上达到了新的无监督性能水平。此外，当将LLM2Vec与监督对比学习相结合时，还在仅使用公开可用数据的模型中实现了最先进的性能。这表明，通过这种简单且有效的方法，原本仅用于生成任务的解码器模型也能被转化为通用的文本编码器，从而在多种NLP任务中表现出色。

方法详解

论文中描述的LLM2Vec方法在代码层面主要涉及以下几个关键的修改，以将decoder-only模型转换为能够生成丰富文本编码的模型：

1. 启用双向注意力：通常，decoder-only模型使用的是单向（因果）注意力机制，这限制了模型只能看到当前标记之前的信息。为了转换这一点，论文中提到通过替换因果注意力掩码（causal attention mask）为全1矩阵，使得每个标记都能看到序列中的所有其他标记，从而实现双向注意力。
2. 蒙版下一个标记预测（MNTP）：这是一个训练目标，结合了下一个标记预测和蒙版语言模型的元素。具体来说，首先在输入序列中随机蒙版一些标记，然后训练模型预测这些蒙版的标记，同时考虑前后文。这种训练方式帮助模型适应其新的双向注意力能力。
3. 无监督对比学习：使用SimCSE（Simple Contrastive Learning of Sentence Embeddings）方法，这种方法通过对同一句子生成两个不同的嵌入表示，并训练模型最大化这两个表示之间的相似度，同时最小化与批次中其他不相关句子表示的相似度。这一步骤不需要配对数据，可以使用任何文本集合。

这些代码修改主要集中在模型的预训练和微调阶段，旨在不仅使模型能够处理更丰富的上下文信息，还提高了模型在不同NLP任务中的通用性和有效性，也就是说我们最终还是需要进行微调训练的，所以下面我们就要展示一些代码来看看如何进行这部分的微调训练。

利用LLM2Vec将Llama 3转化为文本嵌入模型

首先我们安装依赖

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

我们这里演示使用单卡的4090，对于现有模型，我们直接加载现有的模型：

 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")
1
2
3
4
5
6
7
8
9
10

“torch_dtype=torch.bfloat16”是能够在24 GB GPU上运行转换所必需的配置。如果不设置它，模型将是float32参数的原始大小，内存是不够的。

这时，其实这个模型已经可以使用了。但是如果其插入到RAG中。它的性能是不如标准嵌入模型，因为他的运行方式还是因果推断，而不是我们的嵌入。

所以下一步，就需要用MNTP的目标来训练羊驼他。论文的作者还提供了一个脚本:

experiments/run_mntp.py

它目前支持Llama和Mistral架构的模型，所以我们直接可以拿来使用

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

这个脚本需要一个参数，它是JSON格式的配置文件。他们在这里提出了几个例子:

train_configs / mntp

对于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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

该脚本使用bfloat16参数加载模型。将每个设备的批处理大小设置为1，这样可以适合24 GB的GPU。

然后可以开始MNTP训练:

 python llm2vec/experiments/run_mntp.py mntp_config.json
1

使用24gb的4090，或者Google Colab的L4，这需要4天的时间。但是经过MNTP训练后，模型应该会产生更好的结果，特别是对于检索任务。

论文中提到的最后一步是SimCSE，但是作者还没有发布他们的代码，但提到他们会发布的。所以我们直接可以那这个模型来进行测试

使用Llama 3文本嵌入模型

对于训练完成的模型，我们可以直接使用SentenceTransformer加载

 from sentence_transformers import SentenceTransformer
 model = SentenceTransformer("Llama-3-8B-llm2vec-Emb")
1
2

使用LlamaIndex的话可以直接设置

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

这里设置device= ’ cpu '使用cpu加载，这会使RAG系统变慢。可以删除此参数以在GPU上运行它。但是模型是以全精度加载的，所以我们将其加载到CPU上进行测试。因为llm2vec是刚刚发布的，所以还没有任何的量化教程，希望后续会有发布，这样就可以在我们的GPU上完全的使用了。

总结

通过LLM2Vec，我们可以使用LLM作为文本嵌入模型。但是简单地从llm中提取的嵌入模型往往表现不如常规嵌入模型。LLM2Vec的作者提出了新的训练目标MNTP和SimCSE来训练从llm中提取的嵌入模型。这种训练成本很高，但根据作者的说法，可以产生更好的嵌入模型。

论文地址：

https://avoid.overfit.cn/post/67a62b9532b247cc9db87663ce547ff2