为检索增强生成（RAG）微调Embedding模型

简介       

 Embedding模型对于成功的RAG应用至关重要，但它们通常基于通用知识进行训练，这限制了它们在公司或特定领域应用的有效性。为你的特定领域数据定制Embedding可以显著提升你的RAG应用的检索性能。随着Sentence Transformers 3的最新发布，调整Embedding模型比以往任何时候都更加容易。

在这篇博客中，我们将向您展示如何使用2023_10 NVIDIA SEC 报告的合成数据集，为金融RAG应用调整Embedding模型。我们还将利用Matryoshka 表示学习来提升效率。在博客中，我们将进行以下操作：

1、创建并准备嵌入数据集
2、创建基准并评估预训练模型
3、定义使用Matryoshka 表示的损失函数
4、使用SentenceTransformersTrainer调整Embedding模型
5、对比调整后的模型与基准进行评估

Matryoshka Embeddings

Matryoshka表示学习（MRL）是一种旨在创建可以在不显著降低性能的情况下被截断到各种维度的嵌入的技术。这种方法将重要信息前置加载到嵌入的早期维度中，从而在保持高准确度的同时实现下游任务（如检索、分类和聚类）的高效存储和处理。

例如，Matryoshka模型可以在需要存储量减少3倍的情况下保持99.9%的性能。这对于存储和处理资源有限的应用特别有用，如离线应用或大规模检索系统。

注意：此博客旨在在消费者级GPU（如24GB的NVIDIA A10G或RTX 4090/3090）上运行，但可以轻松适应在更大GPU上运行。

在我们开始之前，我们需要安装 Hugging Face 库和 PyTorch，包括 Sentence Transformers, transformers 和 datasets。

# Install Pytorch & other libraries
!pip install "torch==2.1.2" tensorboard
 
# Install Hugging Face libraries
!pip install --upgrade \
  "sentence-transformers>=3" \
  "datasets==2.19.1"  \
  "transformers==4.41.2" 
 

我们将使用Hugging Face Hub作为远程模型版本控制服务。这意味着在训练过程中，我们将自动将我们的模型、日志和信息推送到Hub。您需要在Hugging Face上注册以实现这一点。注册后，我们将使用huggingface_hub包中的login实用工具登录到我们的帐户，并将我们的令牌（访问密钥）存储在磁盘上。

from huggingface_hub import login
 
login(token="", add_to_git_credential=True)  # ADD YOUR TOKEN HERE

1. 创建并准备嵌入数据集

嵌入数据集通常包含文本对（问题，答案/上下文）或表示句子之间关系或相似性的三元组。您选择或可用的数据集格式也将影响可用的损失函数。嵌入数据集的常见格式：

正例对：相关句子的文本对（查询，上下文 | 查询，答案），适用于相似性或语义搜索任务，示例数据集：sentence-transformers/sentence-compression, sentence-transformers/natural-questions。
三元组：由（锚，正例，负例）组成的文本三元组，示例数据集：sentence-transformers/quora-duplicates, nirantk/triplets。
带相似度评分的对：具有表示它们相关程度的相似度评分的句子对，示例数据集：sentence-transformers/stsb, PhilipMay/stsb_multi_mt
更多信息请参阅数据集概览。

我们将使用philschmid/finanical-rag-embedding-dataset，其中包括来自2023_10 NVIDIA SEC 报告的7,000个正例文本对的查询和相应的上下文。

数据集具有以下格式

{"question": "<question>", "context": "<relevant context to answer>"}
{"question": "<question>", "context": "<relevant context to answer>"}
{"question": "<question>", "context": "<relevant context to answer>"}

我们将使用🤗 Datasets库加载我们的开源数据集，并重命名列以匹配sentence-transformers所期望的格式。之后，我们将数据集划分为训练集和测试集，以便能够评估我们的模型。

from datasets import load_dataset
 
# Load dataset from the hub
dataset = load_dataset("philschmid/finanical-rag-embedding-dataset", split="train")
 
# rename columns
dataset = dataset.rename_column("question", "anchor")
dataset = dataset.rename_column("context", "positive")
 
# Add an id column to the dataset
dataset = dataset.add_column("id", range(len(dataset)))
 
# split dataset into a 10% test set
dataset = dataset.train_test_split(test_size=0.1)
 
# save datasets to disk
dataset["train"].to_json("train_dataset.json", orient="records")
dataset["test"].to_json("test_dataset.json", orient="records")

2. 创建基准并评估预训练模型

在创建了我们的数据集之后，我们想要构建一个基准。基准提供了一个参考点，我们可以通过它来衡量自定义模型的性能。通过在特定数据集上评估预训练模型的性能，我们可以了解到模型的初步效果，并确定改进的领域。

对于我们的示例，我们将使用BAAI/bge-base-en-v1.5作为起点。BAAI/bge-base-en-v1.5是大小适中、参数数量仅为109M、隐藏维度为768的最强开源嵌入模型之一，它在MTEB排行榜上取得了63.55分。

我们将使用InformationRetrievalEvaluator来评估模型在给定查询集和文档集上的性能。它将为每个查询检索最相似的k个文档。它测量平均倒序排名（MRR）、@k召回率、平均精确度（MAP）和归一化累积折扣收益（NDCG）。

对我们来说，最重要的是归一化累积折扣收益（NDCG），因为它衡量了排名的质量。它考虑了相关文档在排名中的位置，并对其进行了折扣。折扣值是按对数计算的，这意味着排名较高的相关文档更为重要。

为了进行评估，我们将使用训练集和测试集中的所有“文档”作为潜在检索的来源，以及测试集中的查询。如前所述，我们将使用Matryoshka Representation Learning来提升效率。我们将为以下维度构建基准：64、128、256、512、768。这些正是我们将用于Matryoshka Representation Learning的维度。

import torch
from sentence_transformers import SentenceTransformer
from sentence_transformers.evaluation import (
    InformationRetrievalEvaluator,
    SequentialEvaluator,
)
from sentence_transformers.util import cos_sim
from datasets import load_dataset, concatenate_datasets
 
model_id = "BAAI/bge-base-en-v1.5"  # Hugging Face model ID
matryoshka_dimensions = [768, 512, 256, 128, 64] # Important: large to small
 
# Load a model
model = SentenceTransformer(
    model_id, device="cuda" if torch.cuda.is_available() else "cpu"
)
 
# load test dataset
test_dataset = load_dataset("json", data_files="test_dataset.json", split="train")
train_dataset = load_dataset("json", data_files="train_dataset.json", split="train")
corpus_dataset = concatenate_datasets([train_dataset, test_dataset])
 
# Convert the datasets to dictionaries
corpus = dict(
    zip(corpus_dataset["id"], corpus_dataset["positive"])
)  # Our corpus (cid => document)
queries = dict(
    zip(test_dataset["id"], test_dataset["anchor"])
)  # Our queries (qid => question)
 
# Create a mapping of relevant document (1 in our case) for each query
relevant_docs = {}  # Query ID to relevant documents (qid => set([relevant_cids])
for q_id in queries:
    relevant_docs[q_id] = [q_id]
 
 
matryoshka_evaluators = []
# Iterate over the different dimensions
for dim in matryoshka_dimensions:
    ir_evaluator = InformationRetrievalEvaluator(
        queries=queries,
        corpus=corpus,
        relevant_docs=relevant_docs,
        name=f"dim_{dim}",
        truncate_dim=dim,  # Truncate the embeddings to a certain dimension
        score_functions={"cosine": cos_sim},
    )
    matryoshka_evaluators.append(ir_evaluator)
 
# Create a sequential evaluator
evaluator = SequentialEvaluator(matryoshka_evaluators)

我们将在训练过程中使用评估器进行评估。但首先，让我们创建我们的基准。

# Evaluate the model
results = evaluator(model)
 
# # COMMENT IN for full results
# print(results)
 
# Print the main score
for dim in matryoshka_dimensions:
    key = f"dim_{dim}_cosine_ndcg@10"
    print
    print(f"{key}: {results[key]}")

太好了，使用BAAI/bge-base-en-v1.5作为默认基准，我们得到的结果如下：

维度768：0.7683576219287339
维度512：0.7642500951356254
维度256：0.7546468566985919
维度128：0.7233663127615272
维度64：0.6439999058834552

现在，让我们看看是否可以通过在我们的特定数据集上微调模型来提高这个分数。

3. 为Matryoshka Representation 定义损失函数

在微调嵌入模型时，我们根据数据集格式选择损失函数。对于正文本对，我们可以使用MultipleNegativesRankingLoss与MatryoshkaLoss结合使用。这种方法允许我们利用Matryoshka嵌入的高效性和灵活性，使得可以使用不同的嵌入维度而不会显著影响性能。MultipleNegativesRankingLoss是一个很好的损失函数，如果你只有正对作为输入，它会在损失函数中添加批次负样本，使得每个样本有n-1个负样本。

让我们使用SDPA或Flash Attention 2重新加载我们的模型，并定义一个模型卡片。

from sentence_transformers import SentenceTransformerModelCardData, SentenceTransformer
 
# Hugging Face model ID: https://huggingface.co/BAAI/bge-base-en-v1.5
model_id = "BAAI/bge-base-en-v1.5"
 
# load model with SDPA for using Flash Attention 2
model = SentenceTransformer(
    model_id,
    model_kwargs={"attn_implementation": "sdpa"},
    model_card_data=SentenceTransformerModelCardData(
        language="en",
        license="apache-2.0",
        model_name="BGE base Financial Matryoshka",
    ),
)

在我们加载模型之后，我们可以初始化我们的损失函数。

from sentence_transformers.losses import MatryoshkaLoss, MultipleNegativesRankingLoss
 
matryoshka_dimensions = [768, 512, 256, 128, 64]  # Important: large to small
inner_train_loss = MultipleNegativesRankingLoss(model)
train_loss = MatryoshkaLoss(
    model, inner_train_loss, matryoshka_dims=matryoshka_dimensions
)

4. 使用SentenceTransformersTrainer微调嵌入模型

我们现在准备对模型进行微调。我们将使用SentenceTransformersTrainer，这是transformers库中Trainer的子类，支持所有相同的功能，包括日志记录、评估和检查点。

此外，还有一个SentenceTransformerTrainingArguments类，允许我们指定所有训练参数。

from sentence_transformers import SentenceTransformerTrainingArguments
from sentence_transformers.training_args import BatchSamplers
 
# load train dataset again
train_dataset = load_dataset("json", data_files="train_dataset.json", split="train")
 
# define training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="bge-base-financial-matryoshka", # output directory and hugging face model ID
    num_train_epochs=4,                         # number of epochs
    per_device_train_batch_size=32,             # train batch size
    gradient_accumulation_steps=16,             # for a global batch size of 512
    per_device_eval_batch_size=16,              # evaluation batch size
    warmup_ratio=0.1,                           # warmup ratio
    learning_rate=2e-5,                         # learning rate, 2e-5 is a good value
    lr_scheduler_type="cosine",                 # use constant learning rate scheduler
    optim="adamw_torch_fused",                  # use fused adamw optimizer
    tf32=True,                                  # use tf32 precision
    bf16=True,                                  # use bf16 precision
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # MultipleNegativesRankingLoss benefits from no duplicate samples in a batch
    eval_strategy="epoch",                      # evaluate after each epoch
    save_strategy="epoch",                      # save after each epoch
    logging_steps=10,                           # log every 10 steps
    save_total_limit=3,                         # save only the last 3 models
    load_best_model_at_end=True,                # load the best model when training ends
    metric_for_best_model="eval_dim_128_cosine_ndcg@10",  # Optimizing for the best ndcg@10 score for the 128 dimension
)

我们现在有了创建SentenceTransformersTrainer所需的每一个构建块，接下来就可以开始训练我们的模型了。

from sentence_transformers import SentenceTransformerTrainer
 
trainer = SentenceTransformerTrainer(
    model=model, # bg-base-en-v1
    args=args,  # training arguments
    train_dataset=train_dataset.select_columns(
        ["positive", "anchor"]
    ),  # training dataset
    loss=train_loss,
    evaluator=evaluator,
)

通过调用SentenceTransformerTrainer实例的train()方法开始训练我们的模型。这将启动训练循环，并使我们的模型训练4个epoch。

# start training, the model will be automatically saved to the hub and the output directory
trainer.train()
 
# save the best model
trainer.save_model()
 
# push model to hub
trainer.model.push_to_hub("bge-base-financial-matryoshka")

使用Flash Attention（SDPA）进行4个周期的训练，针对6300个样本，使用g5.2xlarge实例，耗时约00:03:26。每小时实例费用为1.212美元，因此训练总成本仅为0.07美元。

5. 使用基准模型评估微调后的模型

我们在训练过程中对模型进行了评估，但我们也希望在训练结束后与基准模型进行比较。我们使用相同的InformationRetrievalEvaluator来评估模型在给定查询集和语料库集上的性能。

from sentence_transformers import SentenceTransformer
 
fine_tuned_model = SentenceTransformer(
    args.output_dir, device="cuda" if torch.cuda.is_available() else "cpu"
)
# Evaluate the model
results = evaluator(fine_tuned_model)
 
# # COMMENT IN for full results
# print(results)
 
# Print the main score
for dim in matryoshka_dimensions:
    key = f"dim_{dim}_cosine_ndcg@10"
    print(f"{key}: {results[key]}")

我们的微调模型达到：

维度768：0.8253652077429072
维度512：0.8274643684492735
维度256：0.8230326345640168
维度128：0.8184049256658124
维度64：0.7892150398663841

让我们将这些数据放入表格中，并将微调模型的性能与基准模型进行比较。

让我们尝试理解一下：

微调模型在所有维度上都优于基线模型。
BGE基础模型已经非常强大，但在仅使用6300个样本进行微调后，性能仍然提高了约7.4%。
Matryoshka表示学习有效，且在64维度（大小减少12倍）时保持95%的性能，在128维度（大小减少6倍）时保持99%的性能，在256维度（大小减少3倍）时保持超过99.5%的性能。
具有最低维度（64）的微调模型优于具有最高维度（768）的基线模型。
优化128维度的模型允许我们保持与768维度模型99%的性能，同时将存储成本降低6倍并提高余弦相似性搜索。
具有128维度的微调模型在所有维度上都优于具有768维度的基线模型，性能提高了6.51%，同时大小减少了6倍。

结论

嵌入模型对于成功实现RAG应用至关重要，因为如果你无法检索到正确的上下文，就无法生成正确的答案。针对特定领域数据定制嵌入模型，相比于使用通用知识模型，可以显著提高检索性能。微调嵌入模型变得高度可访问，通过使用由LLM生成的合成数据，可以轻松地为特定需求定制模型，从而实现显著的改进。

我们的结果表明，仅使用6.3k样本，微调就可提升性能约7%。训练在消费级GPU上仅需3分钟，通过利用现代技术如Matryoshka表示学习，我们实现了超过99%的性能保留，同时实现了6倍的存储减少和效率提升。