Embedding 模型领域微调实战-优快云博客

本文链接：https://blog.youkuaiyun.com/Python_cocola/article/details/147825559

Embedding 模型领域微调实战

在大模型 RAG 场景中，Embedding 模型多使用开源的 bge/gte 等 Embedding 模型。

这些模型一般在通用数据集上训练，并没有掌握较好的领域知识，从而导致模型在特定领域下的检索效果较差。

本文从代码级别，手把手展示如何对 Embedding 模型进行领域微调，从而提升在特定领域的效果。

0. Embedding 微调简单原理

Embedding 模型通过嵌入将文本转换为固定维度（比如 1024 维）的向量，向量之间的相似度则代表的文本之间的语义相似度。通过微调可以让语义接近的文本的向量相似度变高，语义不接近的文本的向量相似度变低。

其效果就如下图（图源[2]），可简单的理解为让文本簇内集中，簇外分散。

null

1. 准备环境

硬件环境：GTX 4090（可使用其他 NVIDIA 显卡）

软件环境：WSL2 + CUDA 12.4

conda create -n embedding python=3.10 -y
conda activate embedding

# install pytorch with cuda 12.4, see https://pytorch.org/get-started/locally/
# because this bug: https://github.com/huggingface/diffusers/issues/9704, we need to install pytorch-nightly or torch 2.4.
conda install pytorch==2.4.1 torchvision torchaudio pytorch-cuda=12.4 -c pytorch -c nvidia

# install other dependencies
pip install -r requirements.txt

2. 准备数据集

使用标准的 RAG 评估数据集格式作为 Retrieval 评估数据集格式，其中 reference_answers 字段在评估检索效果时可以省略。

评估数据集由如下列组成：

{
    "queries":{
        "<query_id>":"<query>",
        ...
    },
    "corpus":{
        "<corpus_id>":"<corpus>",
        ...
    },
    "relevant_docs":{
        "<query_id>":["<corpus_id>", ...],//每个 Query 可能对应多个 corpus，但是在本案例中，只包含一个。
        ...
    },
    "negative_docs":{
        "<query_id>":["<corpus_id>", ...],
        ...
    },
    "reference_answers":{// 如果只检查检索效果，可以不需要 refrerence_answer 字段。
        "<query_id>":"<reference_answer>",
    }
}

使用 Huggingface 上的 infgrad/retrieval_data_llm 数据集用于训练和验证，其特点是具备挖掘良好的正例和负例，比较适合微调。如下是下载和格式转换命令：

# https://huggingface.co/datasets/infgrad/retrieval_data_llm
PYTHONPATH="." python utils/convert_infgrad_retrieval_data_llm_to_ragdataset.py \
    --dataset "infgrad/retrieval_data_llm" \
    --train_val_split 0.01 \
    --output_path ./data/infgrad_retrieval_data_llm.json
corpus: 369307, train_queries: 182979, val_queries: 1848, save to ./data/infgrad_retrieval_data_llm.json

第二个数据集使用 AirBench-QA-Healthcare-zh，他有 374 条验证集，但是没有训练集，正好可以实践数据合成。如下是下载和格式转换命令：

# https://github.com/AIR-Bench/AIR-Bench/blob/main/docs/available_tasks.md#air-bench_2405
PYTHONPATH="." python utils/convert_airbench_to_ragdataset.py \
    --dataset "qa_healthcare_zh" \
    --train_val_split 1.0 \
    --output_path ./data/airbench_qa_healthcare_zh.json
corpus: 360218, train_queries: 0, val_queries: 374, save to ./data/airbench_qa_healthcare_zh.json

3. 进行基线评估

使用 BAAI/bge-small-zh-v1.5[3] 作为基线模型进行评估。该模型足够小，且容易进行微调。

评估分为2步：

将所有的 corpus 转换为 embedding 向量，并构建 Faiss 索引
进行 eval，在 val 数据上进行评估。

为避免多次操作的时候重复加载模型，制定统一的 embedding cache 机制。请克隆本文对应的代码库后进行操作。

# 在 infgrad 上的指标
PYTHONPATH="." python eval/evaluate_basic.py \
    --dataset_path "./data/infgrad_retrieval_data_llm.json" \
    --encoder "BAAI/bge-small-zh-v1.5" \
    --query_instruction "为这个句子生成表示以用于检索相关文章：" \
    --split"val" \
    --search_top_k 10

{
    "ndcg_at_1": 0.5395,
    "ndcg_at_3": 0.63031,
    "ndcg_at_5": 0.65451,
    "ndcg_at_10": 0.67399,
    "map_at_1": 0.5395,
    "map_at_3": 0.60868,
    "map_at_5": 0.62218,
    "map_at_10": 0.63025,
    "recall_at_1": 0.5395,
    "recall_at_3": 0.69264,
    "recall_at_5": 0.75108,
    "recall_at_10": 0.81115
}

# 在 airbench 上的指标
PYTHONPATH="." python eval/evaluate_basic.py \
    --dataset_path "./data/airbench_qa_healthcare_zh.json" \
    --encoder "BAAI/bge-small-zh-v1.5" \
    --query_instruction "为这个句子生成表示以用于检索相关文章：" \
    --split"val" \
    --search_top_k 10
{
    "ndcg_at_1": 0.40107,
    "ndcg_at_3": 0.32827,
    "ndcg_at_5": 0.30947,
    "ndcg_at_10": 0.32522,
    "map_at_1": 0.0833,
    "map_at_3": 0.14634,
    "map_at_5": 0.18023,
    "map_at_10": 0.21772,
    "recall_at_1": 0.0833,
    "recall_at_3": 0.17145,
    "recall_at_5": 0.23619,
    "recall_at_10": 0.33766
}

这里的 recall、map、ndcg 等指标代表信息检索（IR）领域的指标。在 RAG 场景中，重点关注 recall@N TopN 召回率指标。

4. 微调准备

准备微调数据集和 Loss 函数。根据 sbert loss doc[4] 数据集有三种格式：

仅正样本： (anchor, positive) pairs，一般使用 MultipleNegativesRankingLoss 损失函数
正负样本：(anchor, positive, negative_1, …, negative_n)，一般使用 MultipleNegativesRankingLoss[5] 损失函数
分数样本：(sentence_A, sentence_B, score)，一般使用 CoSENTLoss 损失函数。

第 3 种比较适合做知识蒸馏，也就是用比较强的模型的相似度分数来训练小模型。目前基本使用第一种或第二种。

数据准备和合成可参考如下论文：

null

Wu, T. et al. Towards Robust Text Retrieval with Progressive Learning. Preprint at http://arxiv.org/abs/2311.11691 (2023).

4.1 数据合成

对 airbench_qa_healthcare_zh 数据集实践数据挖掘，从 corpus 中挖掘出正样本（Query和Corpus 的关联关系）。实际生产环境下，corpus 可以通过对文档的识别和切割得到，本文从略。

使用 finetune/data_synthesis.ipynb 进行数据合成。合成后的数据集共有 9810 条QA对。

4.2 难负例挖掘

针对 airbench_qa_healthcare_zh 数据集，使用 finetune/hard_negative_mining.ipynb 进行难负例挖掘。可以控制挖掘的负例的数量，暂定每个 Query 挖掘 15 个负例。

5. 全参数微调

5.1 infgrad 数据集微调

使用 finetune/sft_infgrad.ipynb 进行全参数微调。微调结果保存在 checkpoint/bge-small-zh-v1.5-sft，使用 infgrad_retrieval_data_llm 数据集的正负例数据进行微调，并在 airbench_qa_healthcare_zh 数据集上验证 Held-out 效果。

# 检查在训练的 val 数据上的效果
PYTHONPATH="." python eval/evaluate_basic.py \
    --dataset_path "./data/infgrad_retrieval_data_llm.json" \
    --encoder "checkpoint/bge-small-zh-v1.5-sft" \
    --query_instruction "为这个句子生成表示以用于检索相关文章：" \
    --split"val" \
    --search_top_k 10
{
    "ndcg_at_1": 0.61472,
    "ndcg_at_3": 0.69546,
    "ndcg_at_5": 0.71485,
    "ndcg_at_10": 0.73301,
    "map_at_1": 0.61472,
    "map_at_3": 0.67551,
    "map_at_5": 0.68614,
    "map_at_10": 0.69371,
    "recall_at_1": 0.61472,
    "recall_at_3": 0.75325,
    "recall_at_5": 0.80087,
    "recall_at_10": 0.8566
}

# 检查在非训练数据上的效果
PYTHONPATH="." python eval/evaluate_basic.py \
    --dataset_path "./data/airbench_qa_healthcare_zh.json" \
    --encoder "checkpoint/bge-small-zh-v1.5-sft" \
    --query_instruction "为这个句子生成表示以用于检索相关文章：" \
    --split"val" \
    --search_top_k 10
{
    "ndcg_at_1": 0.37701,
    "ndcg_at_3": 0.31708,
    "ndcg_at_5": 0.29099,
    "ndcg_at_10": 0.30451,
    "map_at_1": 0.07493,
    "map_at_3": 0.14166,
    "map_at_5": 0.16724,
    "map_at_10": 0.19973,
    "recall_at_1": 0.07493,
    "recall_at_3": 0.16802,
    "recall_at_5": 0.22041,
    "recall_at_10": 0.31009
}

可以看到：

微调后的模型在测试集上 recall@3 从 0.63031 提升到 0.75325，提升 12pp。
但是在Held-out （领域外）数据集上 recall@3 从0.17145 降低到 0.16802，降低0.3pp，这就是灾难性遗忘的特点。

5.2 airbench 合成数据集微调

也测试下 sft_airbench.ipynb，使用 airbench_qa_healthcare_zh 我们自己合成的数据集进行微调，并在 infgrad_retrieval_data_llm 数据集上验证 Held-out 效果。微调结果保存在 checkpoint/bge-small-zh-v1.5-sft-airbench

# 检查在训练的 val 数据上的效果
PYTHONPATH="." python eval/evaluate_basic.py \
    --dataset_path "./data/airbench_qa_healthcare_zh.json" \
    --encoder "checkpoint/bge-small-zh-v1.5-sft-airbench" \
    --query_instruction "为这个句子生成表示以用于检索相关文章：" \
    --split"val" \
    --search_top_k 10
{
    "ndcg_at_1": 0.32888,
    "ndcg_at_3": 0.27589,
    "ndcg_at_5": 0.25202,
    "ndcg_at_10": 0.26363,
    "map_at_1": 0.07424,
    "map_at_3": 0.12941,
    "map_at_5": 0.1511,
    "map_at_10": 0.17766,
    "recall_at_1": 0.07424,
    "recall_at_3": 0.14725,
    "recall_at_5": 0.18944,
    "recall_at_10": 0.26303
}

# 检查在非训练数据上的效果
PYTHONPATH="." python eval/evaluate_basic.py \
    --dataset_path "./data/infgrad_retrieval_data_llm.json" \
    --encoder "checkpoint/bge-small-zh-v1.5-sft-airbench" \
    --query_instruction "为这个句子生成表示以用于检索相关文章：" \
    --split"val" \
    --search_top_k 10
{
    "ndcg_at_1": 0.50866,
    "ndcg_at_3": 0.5959,
    "ndcg_at_5": 0.62093,
    "ndcg_at_10": 0.64245,
    "map_at_1": 0.50866,
    "map_at_3": 0.57477,
    "map_at_5": 0.58856,
    "map_at_10": 0.59754,
    "recall_at_1": 0.50866,
    "recall_at_3": 0.65693,
    "recall_at_5": 0.71807,
    "recall_at_10": 0.78409
}

可以看到：

微调后的模型在测试集上 recall@3 从 0.17145 反而降低到 0.14725，降低 2.4pp。这是因为我们合成数据并没有和测试集保持相同的数据分布，反而导致效果下降。证明合成数据在实际生产实践中比较困难。在业界文章中一般用合成数据的测试集进行验证，效果会有提升。
在Held-out （领域外）数据集上 recall@3 从0.69264 降低到 0.65693，降低3.8pp，这也是灾难性遗忘的特点。

6. 结论

如果有较好的领域数据，通过对 Embedding 进行全参微调，可提升模型的领域效果。但是数据合成需要精细化的控制合成的数据分布，否则反而会导致微调后效果下降。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述