Hugging Face Blog

Training and Finetuning Multi-Vector Embedding Models with Sentence Transformers

8.5内容质量

TL;DR · AI 摘要

Sentence Transformers v6.0新增MultiVectorEncoder模型,通过微调可超越通用检索模型。

核心要点

  • MultiVectorEncoder模型使用MaxSim算子实现token级匹配,提升检索精度
  • 在医疗数据集上14.5小时训练的模型超越所有通用检索模型
  • 支持从基础模型和现有模型两种起点进行微调

结构提纲

按章节快速跳转。

  1. 解释多向量模型与传统密集嵌入模型的差异及优势

  2. 详细说明模型、数据集、损失函数等核心训练要素

  3. 演示从基础模型和现有模型进行微调的具体步骤

  4. 展示医疗数据集上14.5小时训练的模型性能对比结果

思维导图

用一张图看清主题之间的关系。

查看大纲文本(无障碍 / 无 JS 友好)
  • 多向量模型训练
    • 模型架构
      • MaxSim算子
      • token级向量
    • 训练流程
      • 数据准备
      • 损失函数
      • 微调方法
    • 评估结果
      • 医疗数据集
      • 性能对比

金句 / Highlights

值得收藏与分享的关键句。

#Sentence Transformers#多向量模型#微调#Hugging Face#ColBERT
打开原文

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

返回文章列表

[-1

]

[0

发布于2026年8月26日

GitHub上的更新

点赞

109

[

  • +103

Tom Aarsen

tomaarsen

关注

Sentence Transformers是一个Python库,用于使用和训练嵌入模型及重排序模型,适用于检索增强生成、语义搜索、语义文本相似性等多种应用场景。其v6.0版本更新引入了第四种模型类型:MultiVectorEncoder,用于ColBERT风格的后期交互检索,并提供了完整的训练方法。在本文中,我将展示如何使用它来微调一个在您的数据上表现优于通用检索器的多向量模型。该方法也可以从头开始训练强大的新多向量模型。以下所有内容均可通过pip install -U "sentence-transformers[train]" 运行。

微调多向量模型涉及多个组件:模型本身、数据集、损失函数、训练参数、评估器以及训练器类。我将逐一查看这些组件,并结合实际示例说明如何使用它们来微调强大的多向量模型。

最后,在评估部分,我将展示我的微调多向量编码器/mLateOn-medical模型,该模型在单个RTX 3090上伴随本文训练14.5小时后,在我的医学检索评估中轻松超越了所有能找到的通用检索模型:包括密集型、稀疏型、词法型和多向量型模型。

如果您对微调密集嵌入模型、稀疏嵌入模型或重排序器更感兴趣,可以阅读我之前的《训练和微调嵌入模型》、《训练和微调稀疏嵌入模型》以及《训练和微调重排序模型》博文。

本文介绍多向量模型的训练。如果想了解如何使用它们,从加载和编码到向量数据库索引,请参阅配套的《使用Sentence Transformers的多向量(后期交互)嵌入模型》博文。

目录

  • 什么是多向量模型?
  • 为什么要微调?
  • 训练组件
  • 微调现有多向量模型 从基础Transformer构建一个模型 应该选择哪个起点?
  • 数据集 Hugging Face Hub上的数据 本地数据 数据集格式
  • 损失函数
  • 训练参数
  • 评估器
  • 训练器回调 多数据集训练
  • 评估 优化索引
  • 致谢
  • 其他资源 训练示例 文档

什么是多向量模型?

密集嵌入模型将整个文本压缩成一个向量,相似性是两个摘要之间的点积。多向量模型(也称为后期交互或ColBERT风格模型)跳过了这种压缩。它为每个标记保留一个小向量,并使用MaxSim运算符将查询与文档进行评分,其中每个查询标记找到其最佳匹配的文档标记,然后将得分相加。标记级匹配保留了单个向量必须平均掉的精细信号,这通常意味着更强的检索效果,但代价是更大的索引。

配套的《多向量嵌入模型》博文详细介绍了架构、编码、评分和索引,因此我将简要介绍本节内容,直接进入训练部分。

为什么要微调?

对多向量模型进行微调可以显著提升其在特定领域中的检索性能:在网页搜索、法律发现、代码搜索和科学文献回顾等不同场景中,词汇、查询风格和相关性概念存在明显差异。由于查询和文档是按标记逐个匹配的,多向量模型能够捕捉到单向量模型容易平均掉的细粒度领域信号,并且对即使是少量的领域内微调数据也能做出良好响应。

此外,大多数公开的检索模型都针对短文本进行了配置。经典ColBERT检查点会在180或300个标记处截断文档,而许多流行的密集模型则在256或512个标记处截断,因为它们的MS MARCO风格训练数据很少超过这个长度。如果您的文档较长,这些模型在评分前会默默丢弃每个文档的大部分内容。在我的医学评估中,平均段落长度为941个标记,我测量到这种截断会导致高达0.24的NDCG@10损失,这个损失远大于任何模型架构之间的差异。当您训练自己的模型时,您可以根据数据需求配置文档长度。

LightOn在代码检索中也遇到了同样的动态,通用的LateOn模型不够用,因此他们训练了LateOn-Code。无论您的领域是医疗、法律、金融还是公司内部文档,都不会有官方模型为您提供支持。这篇博客文章将向您展示如何在单个消费级GPU上,仅用数小时时间就能自行构建模型。

训练组件

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

  • 模型:要微调的模型或要构建的新架构。
  • 数据集:用于训练和评估的数据。
  • 损失函数:衡量模型性能并指导优化过程的函数。
  • 训练参数(可选):影响训练性能、跟踪和调试的参数。
  • 评估器(可选):用于在训练前、中或后评估模型的类。
  • 训练器:将所有训练组件整合在一起。

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

模型

多向量训练为您提供了真正的起点选择,而且这个选择的重要性可能超出您的预期。

微调现有多向量模型

如果您想进一步微调现有多向量模型,完全不需要担心架构问题:

python
from
sentence_transformers
import
MultiVectorEncoder
# 如果内存允许,建议使用fp32进行训练
model = MultiVectorEncoder(
"lightonai/mLateOn-unsupervised"
,
    model_kwargs={
"torch_dtype"
:
"float32"
},
    processor_kwargs={
"model_max_length"
:
8192
},
# tokenizer级别的标记限制
)

检查点自带自己的配置方案:它的查询和文档标记令牌、投影头、评分跳过列表。对于微调,通常需要保留所有这些配置,仅根据数据需求进行调整。首先需要检查的是长度配置,因为许多公开的检查点将文档限制在180到512个标记(参见《为什么需要微调?》),而我的医学段落长度可达1,400个标记。mLateOn系列已经支持主干模型完整的8192标记上下文,但如果初始检查点存在限制,请解除这些限制:

code
# 让模型读取完整文档而不是训练时使用的截断长度
# 例如 GTE-ModernColBERT-v1 默认设置 query_length=48 和 document_length=300
model[
0
].query_length =
None
model[
0
].document_length =
None

在未设置每任务限制的情况下,截断会回退到分词器的 model_max_length 限制,这就是为什么我在上方加载时配置了该限制。

我还做了一个额外的修改,添加了一个标点符号跳过列表,用于排除文档评分和存储中的标点符号标记。在四组消融实验(无、标点符号、停用词、两者都有)中,该方法在质量上略有优势,并且在这些数据上使文档索引体积缩小了9.6%:

code
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 级投影层:

code
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 流水线:一个生成上下文感知 token 嵌入的 Transformer,一个将每个 token 投影到 128 维的 token 级 Dense 层,一个决定评分时哪些 token 有效的 MultiVectorMask,以及一个 token 级 Normalize 层。投影层初始是随机的,因此需要训练才能使模型发挥作用。有趣的是,这种方法也能与强大的密集嵌入主干模型配合使用。在我的实验中,仅使用投影层和 25k 训练对,在 Alibaba-NLP/gte-modernbert-base 上训练的模型,其效果与现有检查点的初始值差距不超过 0.03。

经典的 ColBERT 分词技巧([MASK] 查询扩展、[Q]/[D] 前缀标记、文档长度限制、标点符号跳过列表)默认都是关闭的且可配置。完整的配置选项请参见《创建自定义模型》。就我测试的领域微调而言,在四种配置中测试了 [MASK] 查询扩展,但都没有产生可衡量的差异,因此不必拘泥于经典方案。

你应该选择哪个起点?

在准备这篇博文时,我直接进行了测量,选取了六个起点,每个都使用相同的配方在 MIRIAD 的 25k 医疗问答对上训练,然后在 50k 文档语料库的 1,000 个保留问题上进行评估:

| 起点模型 | 零样本 NDCG@10 | 训练后 NDCG@10 | 差值 | |---------|---------------|---------------|-----| | lightonai/mLateOn-unsupervised | 0.9087 | 0.9398 | +0.0311 | | lightonai/mLateOn | 0.9277 | 0.9319 | +0.0042 | | lightonai/LateOn-unsupervised | 0.9026 | 0.9206 | +0.0180 | | lightonai/LateOn | 0.9185 | 0.9105 | -0.0080 | | lightonai/GTE-ModernColBERT-v1 | 0.9198 | 0.9007 | -0.0191 | | 新建头在 gte-modernbert-base 上 | - | 0.9177 | - |

code

结果让我感到惊讶,并且在两个模型家族中得到了复现。*未监督检查点适应新领域的表现远优于其已完成的兄弟模型,尽管起始表现较低,但最终超越了它们。这些检查点位于大规模对比预训练之后、通用检索的监督微调之前,因此它们保留了所有晚期交互结构,却没有任何通用性调优,而领域训练则需要逆转这些调优。相比之下,已完成的检查点几乎没有变化,甚至在某些学习率下出现了退化。

因此,如果你喜欢的模型家族发布了预监督检查点,请从那里开始。如果没有,使用强检索预训练主干的新投影是一个非常接近的替代方案。尽管从完全完成的检查点继续训练是领域适应中最自然的选择,但却是最弱的选项。

## 数据集

MultiVectorEncoderTrainer 使用 datasets.Dataset 或 datasets.DatasetDict 实例进行训练和评估。你可以从 Hugging Face Datasets Hub 加载数据,也可以使用本地数据(例如 CSV、JSON、Parquet、Arrow 或 SQL 格式)。

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

### Hugging Face Hub 上的数据

你可以使用 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 }) """

code

这就是我在这篇博客文章中将要训练的数据集:来自 MIRIAD 的 440 万个医疗问题,每个问题都与包含其答案的来源段落配对(平均 941 个 token)。像这样的简单(查询,相关段落)对是你为自己的领域收集检索训练数据最容易的方式,正如你将看到的,它们就是你所需要的全部。

### 本地数据

你也可以使用 load_dataset 加载本地常见格式的数据:

from datasets import load_dataset

dataset = load_dataset( "csv" , data_files= "my_file.csv" )

dataset = load_dataset( "json" , data_files= "my_file.json" )

code

如果你的本地数据需要预处理,可以使用 datasets.Dataset.from_dict 通过列表字典初始化数据集:

from datasets import Dataset

queries = [] documents = []

打开文件,执行预处理、过滤、清洗等操作

并将结果追加到列表中

dataset = Dataset.from_dict({ "query" : queries, "document" : documents, })

code

### 数据集格式

确保你的数据集格式与损失函数匹配(或选择与数据集格式匹配的损失函数)非常重要。验证数据集格式是否与损失函数兼容需要两个步骤:

- 如果你的损失函数根据损失函数概览表需要一个 Label,则数据集必须包含名为 "label" 或 "score" 的列。该列将被自动识别为标签。

- 除名为 "label" 或 "score" 的列外,其余所有列均被视为输入,根据损失函数概览表的定义。剩余列的数量必须与所选损失函数的有效输入数量相匹配。这些列的名称无关紧要,仅顺序有影响。

在此基础上还有两项多向量特有的约定:

- 位置查询和文档分配:第一列被嵌入为查询,其余所有列作为文档,与列名无关。此默认行为可通过标准的 router_mapping 训练参数按列覆盖。

- 知识蒸馏格式:每个候选文档占一列,即 (query, document_1, ..., document_N, scores),其中 scores 是每行 N 个教师评分的列表。对于将查询和文档 ID 与独立文本数据集一起存储的知识蒸馏数据集(例如 lightonai/ms-marco-en-bge),可以使用 resolve_ids 实现 ID 到文本的实时解析。

## 损失函数

损失函数用于量化模型在给定数据批次上的表现,使优化器能够更新模型权重以产生更优(即更低)的损失值。适合您任务的损失函数取决于您拥有的数据和要实现的目标。您可以在损失函数概览中找到完整的选项列表。

对于常见问题-答案或问题-段落配对场景,主要使用的是结合 MultiVectorMultipleNegativesRankingLoss 的批次内负样本训练,其中批次中的每个其他文档都作为每个查询的负样本。更大的批次意味着更多的负样本和更强的训练效果,因此在实践中您会希望使用其 GradCache 变体 CachedMultiVectorMultipleNegativesRankingLoss,该变体将有效批次大小与 GPU 容量解耦:

from sentence_transformers import MultiVectorEncoder from sentence_transformers.multi_vector_encoder.losses import CachedMultiVectorMultipleNegativesRankingLoss

model = MultiVectorEncoder( "lightonai/mLateOn-unsupervised" , model_kwargs={ "torch_dtype" : "float32" })

loss = CachedMultiVectorMultipleNegativesRankingLoss( model=model, mini_batch_size= 16 ,

每个块编码的文档数量:限制内存,不影响质量

)

code

mini_batch_size 参数通过以该大小的块编码文档来限制内存,而有效对比批次大小(在我的运行中为 128,且在我的消融实验中更大的批次没有带来进一步提升)则保持自由选择。GradCache 保证无论块大小如何结果都一致,因此对于较小的 GPU 可以仅以时钟时间成本降低该值。当文档长度差异较大时,可以考虑其兄弟参数 mini_batch_num_tokens,该参数将每个块打包到总标记预算而非文档数量,因此异常长文档的块永远不会导致内存激增(我的 mini_batch_size=16 每个文档约 940 个标记对应 mini_batch_num_tokens=15_000)。

一个多向量特有的陷阱是对比损失的默认缩放系数为 scale=1.0,而密集嵌入等价项的默认值为 scale=20.0。这个 20.0 的存在是因为余弦相似度是一个 [-1, 1] 范围内的单一值,对于尖锐的 softmax 来说范围过窄。MaxSim 分数则为每个查询标记计算最佳匹配相似度并求和,因此其范围已扩展到约 [0, 查询长度]:32 个标记的查询最多可得 32 分。因此不要从密集训练脚本复制 scale=20.0,因为这会导致 softmax 饱和并破坏梯度。

对于从更强的教师模型进行知识蒸馏的训练方式(这是目前最强的通用后期交互模型的训练方法),请参阅 MultiVectorDistillKLDivLoss 以及训练概览文档中的知识蒸馏章节。

## 训练参数

你可以通过 MultiVectorEncoderTrainingArguments 类自定义训练过程。该类允许你调整可能影响训练速度并帮助你理解训练过程中发生情况的参数。

有关最有用的训练参数的更多信息,请查看 Multi-Vector Encoder > 训练概览 > 训练参数。为了充分利用训练过程,建议仔细阅读该部分内容。

以下是一个使用我实际训练运行中的参数值的示例:

from sentence_transformers import MultiVectorEncoderTrainingArguments from sentence_transformers.base.sampler import BatchSamplers

args = MultiVectorEncoderTrainingArguments(

必填参数:

output_dir= "models/mLateOn-medical" ,

可选训练参数:

num_train_epochs= 1 , per_device_train_batch_size= 128 ,

由于使用GradCache,有效对比批次大小

per_device_eval_batch_size= 16 , learning_rate= 1e-4 , warmup_steps= 0.05 , prompts={ "question" : "[Q] " , "passage_text" : "[D] " },

以训练列作为键的检查点标记

fp16= False ,

如果你有支持FP16的GPU,请设为True

bf16= True ,

如果你有支持BF16的GPU,请设为True

batch_sampler=BatchSamplers.NO_DUPLICATES,

批内负样本训练需要避免重复

可选跟踪/调试参数:

eval_strategy= "steps" , eval_steps= 0.1 , save_strategy= "steps" , save_steps= 0.05 , logging_steps= 0.01 , run_name= "mLateOn-medical" ,

将用于Trackio、W&B等工具

)

code

其中一些参数需要特别说明:

- prompts:训练过程不会自动应用模型中存储的提示词,因此需要显式地将它们映射到训练列。此处将检查点的[Q]标记映射到问题列,[D]标记映射到段落列,确保训练与推理过程保持一致。

- max_length(故意未设置):该参数仅在训练时限制分词长度,用于在需要比模型完整服务长度更便宜的训练时。我在该数据集上测量了这种快捷方式的代价。在512个token长度下训练,速度约为正常情况的2倍,但NDCG@10指标下降约0.015,且随着数据量增加,这一差距并未缩小,因为模型根本看不到被截断的内容。除非你更需要速度而非质量,否则建议保持未设置,使训练与推理过程一致。

- learning_rate=1e-4:在从5e-6到2e-4的范围内进行扫描后,我发现这个高于常规的学习率效果最好。

## 评估器

为了在训练过程中跟踪模型性能,你可以将 eval_dataset 传递给训练器以获取评估损失,但具体的检索指标能提供更丰富的信息。Sentence Transformers 为多向量模型提供了以下内置评估器:

评估器 | 所需数据
--- | ---

MultiVectorInformationRetrievalEvaluator

| 查询、语料库和相关文档映射

MultiVectorNanoBEIREvaluator

| 无需数据

MultiVectorTripletEvaluator

| (锚点、正样本、负样本)三元组

MultiVectorRerankingEvaluator

| 包含

{'query': '...', 'positive': [...], 'negative': [...]}

的字典列表

MultiVectorDistillationEvaluator

| 包含候选文档和教师评分的查询

对于领域微调,由您自己的保留数据集构建的MultiVectorInformationRetrievalEvaluator才是关键指标。构建时的一个建议是,语料库应足够具有挑战性,以便区分不同模型。在我的案例中,MIRIAD问题是从其自身的来源段落生成的,这使得检索异常容易。仅针对10k个黄金段落进行评估时,几乎所有模型的NDCG@10得分都超过0.97。如果您的评估结果出现这种饱和现象,请添加干扰段落(我使用训练集拆分的去重段落),直到得分分布拉开:

from datasets import load_dataset from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorInformationRetrievalEvaluator

dataset = load_dataset( "tomaarsen/miriad-4.4M-split" )

Gold: 1,000 个评估问题,每个问题对应其自身的段落,使用评估拆分的完整约10k个唯一段落作为初始语料库

corpus = {} queries = {} relevant_docs = {} passage_to_id = {} for idx, row in enumerate (dataset[ "eval" ]): if row[ "passage_text" ] not in passage_to_id: passage_to_id[row[ "passage_text" ]] = f"p { len (passage_to_id)} " corpus[passage_to_id[row[ "passage_text" ]]] = row[ "passage_text" ] if idx < 1_000 : queries[ f"q {idx} " ] = row[ "question" ] relevant_docs[ f"q {idx} " ] = {passage_to_id[row[ "passage_text" ]]}

干扰项:使 haystack 看起来真实的唯一训练段落

seen = set (passage_to_id) for row in dataset[ "train" ]: if len (corpus) >= 200_000 : break if row[ "passage_text" ] not in seen: seen.add(row[ "passage_text" ]) corpus[ f"d { len (corpus)} " ] = row[ "passage_text" ]

evaluator = MultiVectorInformationRetrievalEvaluator( queries=queries, corpus=corpus, relevant_docs=relevant_docs, name= "miriad-dev" , batch_size= 16 , )

results = evaluator(model)

code

## 训练器

MultiVectorEncoderTrainer 是所有先前组件的整合之处。以下是训练 multi-vector-encoder/mLateOn-medical 模型(引言中介绍的模型)的完整脚本:

import logging import string import traceback from datasets import load_dataset from sentence_transformers import ( MultiVectorEncoder, MultiVectorEncoderModelCardData, MultiVectorEncoderTrainer, MultiVectorEncoderTrainingArguments, ) from sentence_transformers.base.sampler import BatchSamplers from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorInformationRetrievalEvaluator from sentence_transformers.multi_vector_encoder.losses import CachedMultiVectorMultipleNegativesRankingLoss

code

logging.basicConfig( format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO )

def main():

1. 加载初始检查点:通过对比预训练,尚未进行监督学习

如果内存允许,建议使用fp32进行训练

model = MultiVectorEncoder( "lightonai/mLateOn-unsupervised", model_kwargs={"torch_dtype": "float32"}, processor_kwargs={"model_max_length": 8192}, model_card_data=MultiVectorEncoderModelCardData( language="en", license="apache-2.0", model_name="mLateOn finetuned on MIRIAD medical retrieval" ) )

2. 解除每任务长度限制,使训练和推理能够处理完整的医学文献段落

model[0].query_length = None model[0].document_length = None

3. 评分时跳过标点符号token:小幅质量提升且索引体积减少9.6%

model[2].skiplist_words = list(string.punctuation) model[2].resolve_with_tokenizer(model.tokenizer)

4. 加载100万组医学问答对

train_dataset = load_dataset( "tomaarsen/miriad-4.4M-split", split="train" ).select(range(1_000_000))

5. 使用GradCache实现的批内负样本:大有效批次,内存受限的分块处理

loss = CachedMultiVectorMultipleNegativesRankingLoss( model=model, mini_batch_size=16 )

6. 轻量开发评估器:用于训练过程监控,使用500个保留问题

对比评估分割的约1万条唯一文献段落。完整的20万协议评估将在之后进行

eval_split = load_dataset( "tomaarsen/miriad-4.4M-split", split="eval" )

corpus, queries, relevant_docs, passage_to_id = {}, {}, {}, {} for idx, row in enumerate(eval_split): if row["passage_text"] not in passage_to_id: passage_to_id[row["passage_text"]] = f"p{len(passage_to_id)}" corpus[passage_to_id[row["passage_text"]]] = row["passage_text"]

if idx < 500: queries[f"q{idx}"] = row["question"] relevant_docs[f"q{idx}"] = {passage_to_id[row["passage_text"]]}

dev_evaluator = MultiVectorInformationRetrievalEvaluator( queries=queries, corpus=corpus, relevant_docs=relevant_docs, name="miriad-dev", batch_size=16 )

7. 训练参数,如上文讨论

run_name = "mLateOn-medical" args = MultiVectorEncoderTrainingArguments( output_dir=f"models/{run_name}", num_train_epochs=1, per_device_train_batch_size=128, per_device_eval_batch_size=16, learning_rate=1e-4, warmup_steps=0.05, prompts={"question": "[Q] ", "passage_text": "[D] "}, fp16=False,

如果拥有支持FP16的GPU请设为True

bf16=True,

如果拥有支持BF16的GPU请设为True

batch_sampler=BatchSamplers.NO_DUPLICATES, eval_strategy="steps", eval_steps=0.1, save_strategy="steps", save_steps=0.05, logging_steps=0.01, run_name=run_name )

8. 创建训练器并开始训练

trainer = MultiVectorEncoderTrainer( model=model, args=args, train_dataset=train_dataset, loss=loss, evaluator=dev_evaluator ) trainer.train()

9. 保存训练后的模型

model.save_pretrained(f"models/{run_name}/final")

10. (可选)推送到Hugging Face Hub

try: model.push_to_hub(run_name) except Exception: logging.error( f"将模型上传到Hugging Face Hub时出错:\n" f"{traceback.format_exc()}" )

code

__name__ == "__main__" : main()

code

这是完整的方案:一个预监督检查点、一百万组领域对、批内负样本、完整文档长度以及高于常规的学习率。该训练在单块RTX 3090显卡上运行了14.5小时,峰值显存占用17.5 GB,所有参数选择都经过实测对比验证而非猜测。

对于预算有限的读者,我的扩展实验表明,使用10万组数据对(75分钟训练时间)的模型性能与完整百万组运行结果的NDCG@10指标差距仅为0.012。其中大部分性能提升集中在训练的第一小时。

### 回调函数

MultiVectorEncoder训练器支持多种transformers.TrainerCallback子类,包括:

- WandbCallback:若已安装wandb,用于将训练指标记录到W&B平台

- TensorBoardCallback:若可访问TensorBoard,用于将训练指标记录到TensorBoard

- CodeCarbonCallback:若已安装codecarbon,用于追踪训练过程的碳排放量

通过training_args中的report_to参数启用这些功能,例如report_to=["wandb", "codecarbon"],需提前安装对应依赖。默认值为"none",设置report_to="all"将激活所有已安装依赖的集成功能。

更多关于这些回调函数及其自定义方法的信息,请参阅Transformers回调函数文档。

### 多数据集训练

通常,表现优异的通用模型会同时在多个数据集上训练。然而,由于各数据集格式差异,这种方案可能面临挑战。幸运的是,MultiVectorEncoderTrainer允许在不强制统一格式的前提下进行多数据集训练,同时支持为不同数据集应用不同的损失函数。以下是同时使用多个数据集进行训练的步骤:

- 使用数据集字典(或datasets.DatasetDict)作为train_dataset(也可选地作为eval_dataset)

- (可选)使用损失函数字典将数据集名称映射到对应损失函数。仅当需要为不同数据集使用不同损失函数时才需要此步骤

每个训练/评估批次将仅包含来自单个数据集的样本。从多个数据集中采样批次的顺序由MultiDatasetBatchSamplers枚举定义,可通过MultiVectorEncoderTrainingArguments的multi_dataset_batch_sampler参数传递。有效选项包括:

- MultiDatasetBatchSamplers.ROUND_ROBIN:按轮询方式从每个数据集中采样,直到某个数据集耗尽。使用此策略时,可能无法用尽每个数据集的所有样本,但会保证各数据集被平等采样

- MultiDatasetBatchSamplers.PROPORTIONAL(默认):按数据集规模比例进行采样。使用此策略时,每个数据集的所有样本都会被使用,且规模更大的数据集会被更频繁地采样

## 评估

为了了解微调模型的表现,我将其与超过50种检索模型配置在MIRIAD评估集上进行了对比测试。该评估集构建方式与上方Evaluator部分完全一致,包含1000个保留的医学问题,在20万个独特段落中搜索(其中1万条黄金段落隐藏在训练集拆分产生的19万条去重干扰项中)。该语料库规模是《Which starting point should you pick?》中5万段落语料库的4倍,因此两个表格之间的得分不具备可比性。

主要结果如下(完整表格见下方可折叠部分):

模型

家族

NDCG@10

multi-vector-encoder/mLateOn-medical(我的模型)

多向量,微调

0.9139

多向量,零样本

0.8520

(能力上限)

0.8502

Qwen/Qwen3-Embedding-4B

密集,零样本

0.7817

voyageai/voyage-4-nano

0.7563

BM25

词法

0.7501

naver/splade-v3

稀疏,零样本

0.6853

微调模型位居榜首,其NDCG@10指标比任何架构最强零样本模型高出+0.062。换句话说,最强零样本模型在75.8%的查询中将正确段落作为首个结果返回,而微调模型在84.9%的查询中实现这一目标,将首名错误率降低了三分之一以上。

架构模式同样清晰,表格前列的模型均为晚期交互架构。在长文档处理中,每个token对应一个向量的方案优于每个文档对应一个向量的方案,即使在训练数据匹配和主干模型匹配的情况下也是如此。DenseOn和LateOn共享训练数据和架构,仅头部不同,晚期交互模型以+0.12的差距胜出,多语言组合(mDenseOn和mLateOn)则以+0.13的差距复现这一结果。规模优势也无法拯救单向量模型。Qwen3-Embedding-4B作为最强的密集模型,其活跃参数(非嵌入参数)数量约为我的33倍,但仍与最优值相差0.13,8B版本的得分甚至低于4B版本。

BM25的表现也令人意外,其表现优于所有稀疏模型、所有截断限制的多向量模型,以及除三个密集模型外的所有模型:多十亿参数的Qwen3-Embedding-4B和8B,以及voyage-4-nano(该模型通过完整读取32k token上下文仅以0.006的微弱优势胜出)。但不要期待这种优势能迁移到你的数据集。MIRIAD的问题是从段落生成的,因此查询与其黄金段落之间的词汇重叠度远高于典型检索场景,而BM25的无限上下文长度使其能够利用所有重叠词汇,而大多数神经网络模型都会进行截断。运行一个BM25基线模型成本低廉且始终值得,但不要指望能获得这种幅度的提升。

按得分排序并按架构家族着色的完整模型概览。

点击查看完整评估表格

acc@1

0.849

多向量

0.758

@1024

0.763

0.8485

0.760

0.8304

0.733

mixedbread-ai/mxbai-edge-colbert-v0-32m

0.8186

0.727

密集

0.669

Qwen/Qwen3-Embedding-8B

0.7747

0.654

perplexity-ai/pplx-embed-v1-late-0.6b

0.7702

0.632

lightonai/ColBERT-Zero

0.7613

0.675

LiquidAI/LFM2.5-ColBERT-350M

0.7582

0.664

0.638

0.641

jinaai/jina-embeddings-v5-text-small-retrieval

0.7470

0.620

Qwen/Qwen3-Embedding-0.6B

0.7408

perplexity-ai/pplx-embed-v1-0.6b

0.7384

0.615

0.7350

0.639

mixedbread-ai/mxbai-edge-colbert-v0-17m

0.7271

0.631

answerdotai/answerai-colbert-small-v1

@512

0.7264

lightonai/DenseOn

0.7239

0.597

lightonai/mDenseOn

0.7227

0.585

jinaai/jina-embeddings-v5-text-nano-retrieval

0.7206

0.587

microsoft/harrier-oss-v1-0.6b

0.7126

0.572

Alibaba-NLP/gte-modernbert-base

0.7102

0.582

Snowflake/snowflake-arctic-embed-l-v2.0

0.7068

0.568

0.7008

0.570

google/embeddinggemma-300m

0.7000

0.563

0.6943

稀疏

0.574

ibm-granite/granite-embedding-small-english-r2

0.6813

0.546

naver/splade-v3-distilbert

0.6806

0.567

codefuse-ai/F2LLM-v2-0.6B

0.6799

0.536

colbert-ir/colbertv2.0

0.6785

0.571

prithivida/Splade_PP_en_v1

0.6755

0.577

0.6713

0.561

tomaarsen/embeddinggemma-300m-miriad-unsloth

密集,微调

0.6705

0.530

lightonai/LateOn-regularized

0.6673

0.554

0.6672

0.553

0.6612

0.555

opensearch-project/opensearch-neural-sparse-encoding-v2-distill

0.6518

0.531

nomic-ai/nomic-embed-text-v1.5

(prompted)

0.6387

0.498

mixedbread-ai/mxbai-embed-large-v1

0.6355

0.502

BAAI/bge-large-en-v1.5

0.6308

jinaai/jina-colbert-v2

0.6218

0.504

0.6203

0.487

0.6184

0.514

tomaarsen/splade-modernbert-base-miriad

稀疏,微调

0.6142

0.473

NeuML/biomedbert-base-colbert

0.5963

0.463

BAAI/bge-base-en-v1.5

0.5930

0.454

BAAI/bge-small-en-v1.5

0.5881

0.457

sentence-transformers/all-mpnet-base-v2

0.5159

0.396

0.4992

0.401

mixedbread-ai/mxbai-colbert-large-v1

0.4690

0.358

sentence-transformers/static-retrieval-mrl-en-v1

0.4614

0.323

sentence-transformers/all-MiniLM-L6-v2

0.4458

0.321

0.4347

0.346

标记 @N 的模型在评估时已解除其文档长度上限至 N 个 token,因为这些模型原本的长度限制(180 至 512 个 token)会截断平均 941 个 token 的段落。对于所有多向量模型,解除长度限制使 NDCG@10 提升了 +0.08 至 +0.24,而即使是密集型 DenseOn 模型也获得了 +0.03 的提升。

请注意,这并不意味着 multi-vector-encoder/mLateOn-medical 在所有领域都是最强的模型。它只是在我的领域中表现最强。这完全没问题,因为只要这个模型在我的数据上表现良好即可。

不要低估在你的领域上对多向量模型进行微调的威力。在单块消费级 GPU 上耗费十四又二分之一小时训练出的模型,在这些数据上表现远超任何通用检索器,而且训练方法仅需一个脚本,无需教师模型,也无需挖掘负样本!

### 优化索引

对多向量检索的合理质疑是索引大小,而本领域几乎是最糟糕的情况。每个 token 存储一个向量,我的模型每个段落需要约 878 个向量,因此包含 20 万个段落的语料库在 fp16 格式下需要约 45 GB,而密集型模型所需空间远低于 1 GB。文档长度正是造成这种差距的原因。配套文章中的 Natural Questions 段落平均每个段落约有 125 个 token 向量,仅为本案例的七分之一,因此短段落语料库的索引起点远小于本案例。HierarchicalTokenPooling 模块通过聚类每个文档的 token 向量并存储聚类中心,将向量数量压缩到原来的 1/pool_factor:

from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

pooling = HierarchicalTokenPooling(pool_factor= 4 ) document_embeddings = model.encode_document(passages, token_pooling=pooling)

code

我在最终模型上进行后验测试,未使用任何感知池化的训练,结果显示在长文档上这种压缩方式成本极低。

实心点表示未压缩的嵌入向量,这样每个模型家族都能以相同方式计数并进行精确搜索。不过你不会以这种方式部署它们。密集型索引通常使用 int8 或二进制量化配合重排序,稀疏索引会压缩倒排索引,多向量索引则使用 PLAID 风格的残差压缩。不要将这些点理解为你需要购买的硬盘容量,而是相对存储成本。

向量池化(token pooling)对应实线。向量数量减半仅导致 NDCG@10 下降 0.0033,且不影响第一名准确率,仅保留四分之一向量(11.2 GB)仍能获得 0.8991 的得分。曲线还会继续(我测试到向量数量的十分之一时得分仍为 0.8765),但量化技术可用后,进一步压缩池化的理由就非常有限了,这正是下方虚线所代表的内容。

虚线表示实际部署可能呈现的效果。我提前将模型和基准测试提供给 Omar Khattab,他使用 fast-plaid 在 1 位残差量化下测量了这些配置,采用紧凑的 17 位中心点 ID 和 18 位文档 ID 替代常规的 64 位整数,并结合文档侧剪枝:

配置

保留向量比例

索引大小

1-bit PLAID(全部向量)

100%

3.37 GB

0.8984

1-bit PLAID + 剪枝

65%

2.23 GB

0.8830

42%

1.45 GB

0.8642

第一行的索引体积仅为原始嵌入向量的 1/13,NDCG@10 为 0.0155。这比任何池化曲线上的表现都更优。量化压缩每个向量,而池化和剪枝减少保留数量,二者效果叠加,其中量化应作为首选优化手段。进一步压缩后最后一行达到 1.45 GB,体积甚至小于 Qwen3-Embedding-8B 的 fp16 嵌入(1.64 GB),同时得分高出 0.0895。关于多向量索引体积过大的质疑,在正确配置的索引面前不攻自破。

此处的剪枝方法较为基础,仅用于验证在量化基础上的向量削减有效性,因此底部两行数据应视为下限而非前沿水平。若希望完全避免手动调整量化参数,配套文章的「索引构建」章节已涵盖 fast-plaid、Qdrant、Weaviate 和 Vespa 的实现方案。

多向量检索的代价完全取决于索引规模。该语料库的原始嵌入体积为 45 GB,而正确配置的索引体积可压缩至原大小的 1/7 且精度几乎无损。索引的优化程度值得与模型检查点同等重视。

## 致谢

感谢 Omar Khattab 在《优化索引》一文中测量量化剪枝索引配置,并参与关于晚期交互索引成本的讨论。

## 附加资源

### 训练示例

这些页面包含训练示例、解释说明及训练脚本链接,可帮助您熟悉多向量训练流程:

- MIRIAD:医疗检索领域的专项训练,是本文方法的早期简化版
- MS MARCO:对比学习与知识蒸馏方案
- 多模态:ColPali 风格的视觉文档检索训练
- PEFT 适配器:基于 LoRA 的参数高效微调

### 文档资料

如需进一步学习,建议查阅 Sentence Transformers 相关资源:

- 安装指南
- 快速入门
- 使用说明
- 自定义模型构建
- 预训练模型
- 训练概览(本文是对该文档的浓缩提炼)
- 损失函数概览
- API 参考文档

以下高级内容可能也值得关注:

- 分布式训练
- Sentence Transformers 多向量(晚期交互)嵌入模型使用指南

## 本文提及的模型 10

## 本文提及的数据集 3

更多博客文章

nlp

guide

community

## 使用 Sentence Transformers 的多向量(晚期交互)嵌入模型

104

2026年8月18日

open-source

Hot

## 欢迎使用 EmbeddingGemma,Google 新型高效嵌入模型

- +2

277

2025年9月4日

### 社区

lucilamorgan

2天前

[2

> 微调多向量模型涉及多个组件:模型本身、数据集、损失函数、训练参数、评估器以及训练器类。我将逐一探讨这些组件,并结合实际示例说明它们如何用于微调强大的多向量模型。

这是一篇关于多向量嵌入模型的深入分析... 🔥 Sentence Transformers v6.0 中引入的 MultiVectorEncoder 特别适用于 ColBERT 风格的晚期交互检索,这种设计通过保留多个表示形式,能够比单一密集嵌入捕获更细粒度的匹配信息。该框架同时支持 RAG、语义搜索和重排序,这使它成为现代检索系统的非常实用的方向。🚀

我也很喜欢文章将微调过程分解为各个独立组件的方式... 数据集、损失函数、训练参数、评估器和训练器类。🧠 每个组件都有实际示例,这使得方法更易于理解和复现,特别是对于那些希望超越通用检索器、为自己的领域优化模型的人而言。

医疗检索结果可能是最令人印象深刻的部分... 💡 仅在单个 RTX 3090 上训练 14.5 小时就获得比密集型、稀疏型、词法型和其他多向量方法更好的效果,这充分证明了领域特定微调的价值。这是一个很好的提醒:检索性能的最大提升往往来自于将模型和训练数据紧密匹配实际搜索任务,而不是简单选择最新的通用模型。👏📚

总体而言,这是一份对任何尝试高级检索流水线的人来说都非常实用的资源... ⚡ 从头训练强大多向量模型以及微调现有模型的能力,为专用搜索和 RAG 应用打开了许多可能性。我特别欣赏整个工作流程可以使用 sentence-transformers[train] 设置... 这使得原本相对复杂的检索训练过程显得更加易于上手。👍🤖

See translation

Reply

Edit

Preview

Upload images, audio, and videos by dragging in the text input, pasting, or

clicking here

.

Tap or paste here to upload images

Comment

· Sign up or log in to comment

- +97
 /think