使用Sentence Transformers训练与微调多向量嵌入模型

来源：https://huggingface.co/blog/train-multi-vector-encoder Sentence Transformers (https://sbert.net/) 是一个Python库，用于使用和训练适用于广泛场景的嵌入模型与重排序模型，例如检索增强生成、语义搜索、语义文本相似性等。其6.0版本更新引入了第四种模型类型：MultiVectorEncoder，用于ColBERT风格的延迟交互检索，并提供了完整的训练方法。在本文中，我将展示如何使用它来微调一个多向量模型，使其在您的数据上超越通用检索器。这种方法也可用于从头开始训练强大的新多向量模型。以下所有内容均可通过运行pip install -U "sentence-transformers[train]"实现。

微调多向量模型涉及多个组件：模型本身、数据集、损失函数、训练参数、评估器以及训练器类。我将逐一介绍这些组件，并附上如何使用它们微调强大多向量模型的实用示例。

最后，在评估部分，我将展示我微调的multi-vector-encoder/mLateOn-medical (https://huggingface.co/multi-vector-encoder/mLateOn-medical) 模型——该模型在单张RTX 3090上训练了14.5小时——在我的医疗检索评估中轻松超越了所有我能找到的通用检索模型，无论是稠密、稀疏、词法还是多向量模型。

在MIRIAD上的NDCG@10与活动参数量对比：微调后的mLateOn-medical以远小于最强通用模型(https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/train-multi-vector-encoder/mve_medical_model_size_ndcg.png)的参数规模达到了最佳性能。

如果您有兴趣微调稠密嵌入模型、稀疏嵌入模型或重排序模型，请参阅我之前的训练与微调嵌入模型 (https://huggingface.co/blog/train-sentence-transformers)、训练与微调稀疏嵌入模型 (https://huggingface.co/blog/train-sparse-encoder) 和训练与微调重排序模型 (https://huggingface.co/blog/train-reranker) 博客文章。

本文介绍训练多向量模型。如果您想了解如何使用它们，从加载、编码到向量数据库中的索引，请参阅配套的使用Sentence Transformers的多向量（延迟交互）嵌入模型 (https://huggingface.co/blog/multi-vector-encoder) 博客文章。

目录

• 什么是多向量模型？ (https://huggingface.co/blog/train-multi-vector-encoder#what-are-multi-vector-models)
• 为什么要微调？ (https://huggingface.co/blog/train-multi-vector-encoder#why-finetune)
• 训练组件 (https://huggingface.co/blog/train-multi-vector-encoder#training-components)
• 模型 (https://huggingface.co/blog/train-multi-vector-encoder#model)
  • 微调现有的多向量模型 (https://huggingface.co/blog/train-multi-vector-encoder#finetuning-an-existing-multi-vector-model)
  • 从基础Transformer构建模型 (https://huggingface.co/blog/train-multi-vector-encoder#building-one-from-a-base-transformer)
  • 应该选择哪个起点？ (https://huggingface.co/blog/train-multi-vector-encoder#which-starting-point-should-you-pick)
• 数据集 (https://huggingface.co/blog/train-multi-vector-encoder#dataset)
  • Hugging Face Hub上的数据 (https://huggingface.co/blog/train-multi-vector-encoder#data-on-the-hugging-face-hub)
  • 本地数据 (https://huggingface.co/blog/train-multi-vector-encoder#local-data)
  • 数据集格式 (https://huggingface.co/blog/train-multi-vector-encoder#dataset-format)
• 损失函数 (https://huggingface.co/blog/train-multi-vector-encoder#loss-function)
• 训练参数 (https://huggingface.co/blog/train-multi-vector-encoder#training-arguments)
• 评估器 (https://huggingface.co/blog/train-multi-vector-encoder#evaluator)
• 训练器 (https://huggingface.co/blog/train-multi-vector-encoder#trainer)
  • 回调 (https://huggingface.co/blog/train-multi-vector-encoder#callbacks)
  • 多数据集训练 (https://huggingface.co/blog/train-multi-vector-encoder#multi-dataset-training)
• 评估 (https://huggingface.co/blog/train-multi-vector-encoder#evaluation)
  • 优化索引 (https://huggingface.co/blog/train-multi-vector-encoder#optimizing-the-index)
• 致谢 (https://huggingface.co/blog/train-multi-vector-encoder#acknowledgements)
• 附加资源 (https://huggingface.co/blog/train-multi-vector-encoder#additional-resources)
  • 训练示例 (https://huggingface.co/blog/train-multi-vector-encoder#training-examples)
  • 文档 (https://huggingface.co/blog/train-multi-vector-encoder#documentation)

什么是多向量模型？

稠密嵌入模型将整个文本压缩成一个向量，相似度是两个此类摘要向量之间的点积。多向量模型（也称为延迟交互或ColBERT风格模型）跳过了这种压缩。它保留每个token对应一个小向量，并使用MaxSim算子对查询与文档进行评分，其中每个查询token找到其最匹配的文档token，并将分数求和。token级别的匹配精确地保留了单向量模型不得不平均掉的细粒度信号，这通常意味着更强的检索能力，代价是更大的索引。

配套的多向量嵌入模型 (https://huggingface.co/blog/multi-vector-encoder) 博客文章详细介绍了架构、编码、评分和索引，因此我将保持本节简短，直接进入训练部分。

稠密嵌入与多向量延迟交互对比 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/multi-vector-encoder/maxsim_explainer.gif)

为什么要微调？

微调多向量模型可以显著提升其在特定领域数据上的检索性能：词汇、查询风格以及相关性概念在网络搜索、法律发现、代码搜索和科学文献综述之间都存在差异。由于查询和文档是逐token匹配的，多向量模型能够捕捉到单向量模型容易平均掉的细粒度领域信号，并且它们对少量领域内微调数据就能有很好的响应。

除此之外，大多数已发布的检索模型是为短文本段落配置的。经典的ColBERT检查点在180或300个token处截断文档，许多流行的稠密模型在256或512个token处截断，因为它们的MS MARCO风格训练数据很少超过这个长度。如果您的文档较长，这些模型会在评分前静默丢弃大部分文档内容。在我的医疗评估中（段落平均长度为941个token），我测得这种截断最多会损失0.24个NDCG@10，这比任何模型架构之间的差异都要大得多。当您训练自己的模型时，可以配置您的数据所需的文档长度。

LightOn在代码检索中也遇到了同样的情况，通用的LateOn (https://huggingface.co/lightonai/LateOn) 不够用，他们训练了LateOn-Code (https://huggingface.co/lightonai/LateOn-Code)。您的领域，无论是医疗、法律、金融还是公司内部文档，都不会有官方模型。本文将向您展示如何在几小时内，在单张消费级GPU上自己构建它。

训练组件

训练MultiVectorEncoder模型涉及以下组件：

1. 模型 (https://huggingface.co/blog/train-multi-vector-encoder#model)：要微调的模型或要从头构建的架构。
2. 数据集 (https://huggingface.co/blog/train-multi-vector-encoder#dataset)：用于训练和评估的数据。
3. 损失函数 (https://huggingface.co/blog/train-multi-vector-encoder#loss-function)：衡量模型性能并指导优化过程的函数。
4. 训练参数 (https://huggingface.co/blog/train-multi-vector-encoder#training-arguments)（可选）：影响训练性能、跟踪和调试的参数。
5. 评估器 (https://huggingface.co/blog/train-multi-vector-encoder#evaluator)（可选）：用于在训练前、训练中或训练后评估模型的类。
6. 训练器 (https://huggingface.co/blog/train-multi-vector-encoder#trainer)：整合所有训练组件。

让我们更详细地了解每个组件。

模型

多向量训练为您提供了真正的起点选择，其重要性可能超出您的预期。

微调现有的多向量模型

如果您想进一步微调现有的多向量模型，则完全不必担心架构问题：

from sentence_transformers import MultiVectorEncoder

# 如果内存允许，加载fp32格式更利于训练
model = MultiVectorEncoder(
    "lightonai/mLateOn-unsupervised",
    model_kwargs={"torch_dtype": "float32"},
    processor_kwargs={"model_max_length": 8192},  # 分词器级别的token限制
)

该检查点带来了自己的配置：它的查询和文档标记token、它的投影头、它的评分管线列表。对于微调，您通常希望保留所有这些，只改变您的数据所要求的内容。首先检查的是长度配置，因为许多已发布的检查点将文档限制在180到512个token（参见为什么要微调？ (https://huggingface.co/blog/train-multi-vector-encoder#why-finetune)），而我的医疗段落长度可达1,400个token。mLateOn系列已经支持主干网络完整的8192个token上下文，但如果您的起点检查点带有上限，请解除它们：

# 让模型阅读完整文档，而不是其训练时的上限，
# 例如 GTE-ModernColBERT-v1 出厂设置了 query_length=48 和 document_length=300
model[0].query_length = None
model[0].document_length = None

由于未设置每个任务的上限，截断将回退到分词器的model_max_length，这就是为什么我在上面加载时配置了该限制。

我做了另一个更改，添加了一个标点符号跳过列表，从文档端的评分和存储中排除标点符号token。在一个4种设置的消融实验中（无、标点符号、停用词、两者），它在质量上略有优势，并且在此数据上免费将文档索引大小缩小了9.6%：

import string

# model[2] 是 MultiVectorMask 模块
model[2].skiplist_words = list(string.punctuation)
model[2].resolve_with_tokenizer(model.tokenizer)  # token ID已被缓存，更改后需重新解析

从基础Transformer构建模型

您也可以将MultiVectorEncoder指向任何基础Transformer，系统会为您附加一个全新的、随机初始化的token级投影头：

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("answerdotai/ModernBERT-base", model_kwargs={"torch_dtype": "float32"})
# MultiVectorEncoder(
#   (0): Transformer({..., 'architecture': 'ModernBertModel'})
#   (1): Dense({'in_features': 768, 'out_features': 128, 'bias': False, ...})
#   (2): MultiVectorMask({'skiplist_words': [], 'skiplist_tasks': ['document'], ...})
#   (3): Normalize({...})
# )

这就是经典的ColBERT流水线：一个Transformer生成上下文化的token嵌入，一个token级的Dense将每个嵌入降维到128维，一个MultiVectorMask决定哪些token在评分时计入，以及一个token级的Normalize。投影头是随机初始化的，因此需要训练才能使该模型有用。有趣的是，这在使用强大的稠密嵌入主干网络时也有效。在我的实验中，在Alibaba-NLP/gte-modernbert-base (https://huggingface.co/Alibaba-NLP/gte-modernbert-base)上随机初始化的投影头，仅通过投影头和25k训练对，其性能就达到了现有检查点起点的0.03以内。

经典的ColBERT分词技巧（[MASK]查询扩展、[Q]/[D]前缀token、文档长度上限、标点符号跳过列表）默认都是关闭的，可配置。完整选项请参阅创建自定义模型 (https://sbert.net/docs/multi_vector_encoder/usage/custom_models.html)。就我测试的情况来看，我在领域微调中测试了四种配置下的[MASK]查询扩展，没有一种产生可衡量的差异，因此不必拘泥于经典配方。

应该选择哪个起点？

在准备本文时，我直接测量了这一点，选择了六个起点，使用完全相同的配方在MIRIAD (https://huggingface.co/datasets/tomaarsen/miriad-4.4M-split) 的25k医疗问题-段落对上进行训练，然后在50,000个段落语料库上对1,000个预留问题进行评估：

结果让我感到惊讶，并且在两个模型系列中都得到了复现。*-unsupervised检查点对新领域的适应能力远胜于其精调后的同类模型，尽管起点更低，但最终超越了它们。这些检查点位于大规模对比预训练之后、通用检索监督微调之前，因此它们保留了所有延迟交互结构，而没有需要领域训练去消除的通用微调。相反，精调后的检查点在我尝试的所有学习率下几乎原地踏步，甚至有所退步。

因此，如果您喜欢的模型系列发布了监督预训练前的检查点，请从那里开始。如果不是，在强大的检索预训练主干上使用新的投影头是次优选择。从完全精调后的检查点继续进行领域适应是最弱的选择，尽管它感觉最自然。

数据集

MultiVectorEncoderTrainer (https://sbert.net/docs/package_reference/multi_vector_encoder/trainer.html)使用datasets.Dataset (https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset)或datasets.DatasetDict (https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.DatasetDict)实例进行训练和评估。您可以从Hugging Face Datasets Hub (https://huggingface.co/datasets)加载数据，或使用您喜欢的任何格式的本地数据（例如CSV、JSON、Parquet、Arrow或SQL）。

注意： 许多与Sentence Transformers开箱即用的公开数据集已在Hugging Face Hub上标记了sentence-transformers，因此您可以在https://huggingface.co/datasets?other=sentence-transformers轻松找到它们。建议浏览这些数据集，寻找可能对您的任务、领域或语言有用的数据。

Hugging Face Hub上的数据

您可以使用load_dataset (https://huggingface.co/docs/datasets/main/en/package_reference/loading_methods#datasets.load_dataset)函数从Hub上的数据集加载数据：

from datasets import load_dataset

train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train")

print(train_dataset)
"""
Dataset({
    features: ['question', 'passage_text'],
    num_rows: 4467542
})
"""

这是从Hub加载数据的最简单方式。在