← 文章 / AI技术
HuggingFace博客 10小时前 · 2026-09-02 05:02:09 · 0 阅读

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

Sentence Transformers 是一个 Python 库,用于使用和训练面向各种应用的嵌入与重排序模型,包括检索增强生成、语义搜索、语义文本相似度等场景。其 v6.0 版本新增了第四种模型类型:用于 ColBERT 式晚期交互检索的 MultiVectorEncoder,并配套提供了一整套训练方案。在本文中,我会演示如何利用它微调多向量模型,让它在你的数据上超越通用检索器。该方法也能从零训练出强大的多向量模型。下面所有内容只需运行 pip install -U "sentence-transformers[train]" 即可上手。

微调多向量模型涉及多个组件:模型本身、数据集、损失函数、训练参数、评估器以及训练器类。我会逐一介绍每个组件,并配合实际示例展示如何用它们微调出强大的多向量模型。

最后,在评估章节,我会展示我在撰写本文的同时用单张 RTX 3090 耗时 14.5 小时微调出的 multi-vector-encoder/mLateOn-medical 模型,在我的医学检索评估任务上轻松超越了能找到的所有通用检索模型——无论是稠密的、稀疏的、词法层面的,还是多向量的。

在 MIRIAD 上以 NDCG@10 与活跃参数量对比:微调后的 mLateOn-medical 以远低于最强通用模型的参数量登顶

如果你感兴趣的是微调稠密嵌入模型、稀疏嵌入模型或重排序器,可以阅读我之前的几篇博客:《训练与微调嵌入模型》《训练与微调稀疏嵌入模型》 以及 《训练与微调重排序模型》

本文讲的是如何训练多向量模型。如果你想了解如何使用它们——从加载、编码到在向量数据库中建索引——可以参考配套博客 Multi-Vector (Late Interaction) Embedding Models with Sentence Transformers

目录

什么是多向量模型?

稠密嵌入模型会把整段文本压缩成一个向量,两段文本之间的相似度就是这两个向量的点积。多向量模型(也叫晚期交互模型或 ColBERT 风格模型)跳过了这一步压缩:它为每个 token 保留一个小向量,用 MaxSim 算子对查询和文档打分——让每个查询 token 找到与其最匹配的文档 token,再把所有得分加起来。token 级的匹配保留了那些被单向量模型不得不"平均掉"的细粒度信号,通常意味着更强的检索效果,代价则是更大的索引体积。

配套的多向量嵌入模型文章详细介绍了架构、编码、打分和索引,所以这一节我就简短带过,直接进入训练部分。

稠密嵌入与多向量晚期交互对比

为什么要微调?

微调多向量模型能显著提升它在特定领域上的检索效果:不同场景下——网页搜索、法律案件检索、代码搜索、学术文献综述——用词、查询风格、相关性的定义都不一样。由于查询和文档是逐 token 比对的,多向量模型能够捕捉到单向量模型容易忽略的细粒度领域信号,因此即便只用少量领域内数据进行微调,效果提升也会非常明显。

除此之外,市面上大多数已发布的检索模型都是为短文本配置的。经典的 ColBERT 权重把文档截断到 180 或 300 token,很多主流稠密模型则截断到 256 或 512 token——因为它们的 MS MARCO 风格训练数据几乎不会超过这个长度。如果你的文档很长,这些模型会在打分之前就把每个文档的大部分内容悄悄丢弃。我在一次医学评估中测试过,文档平均长度为 941 token,这种截断最多会损失 0.24 的 NDCG@10,远超不同模型架构之间的差异。当你训练自己的模型时,就可以根据你的数据配置合适的文档长度。

LightOn 在代码检索场景中也遇到了同样的问题——通用的 LateOn 并不够用,于是他们专门训练了 LateOn-Code。无论你的领域是医学、法律、金融,还是公司内部的文档,都不会有现成的官方模型可以直接用。本文将向你展示如何在几小时内,仅用一块消费级 GPU 自己打造一个合适的模型。

训练组件

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

  1. Model(模型):待微调的模型或要从零搭建的架构。
  2. Dataset(数据集):用于训练和评估的数据。
  3. Loss Function(损失函数):衡量模型表现并指导优化过程的函数。
  4. Training Arguments(训练参数)(可选):影响训练性能、日志记录和调试的参数。
  5. Evaluator(评估器)(可选):用于在训练前、训练中或训练后评估模型的类。
  6. Trainer(训练器):将所有训练组件整合在一起。

下面我们来逐一了解每个组件。

Model

多向量训练让你在起点的选择上有充分的余地,而这一点比想象中更重要。

微调已有的多向量模型

如果你想进一步微调一个已有的多向量模型,完全不需要操心架构的问题:

from sentence_transformers import MultiVectorEncoder

# 训练时如果显存允许,建议使用 fp32 加载
model = MultiVectorEncoder(
    "lightonai/mLateOn-unsupervised",
    model_kwargs={"torch_dtype": "float32"},
    processor_kwargs={"model_max_length": 8192},  # tokenizer 层面的 token 上限
)

这个 checkpoint 自带一套训练配置,包括 query 和 document 的标记 token、projection head 以及 scoring skiplist。微调时,通常应保留这些组件,只根据数据需求修改必要的部分。首先要检查长度配置,因为许多已发布的 checkpoint 会将文档长度限制在 180 到 512 tokens(见为何微调?),而我的医学段落长达 1,400 tokens。mLateOn 系列已经支持 backbone 的完整 8,192 tokens 上下文;但如果初始 checkpoint 存在长度上限,就将其解除:

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

取消针对各任务的限制后,截断处理将回退到 tokenizer 的 model_max_length,因此我会在上面的加载阶段配置这一上限。

我还做了一处调整:新增标点 skiplist,在 document 端的评分和存储中排除标点 token。在四组消融实验(不排除、排除标点、排除停用词、同时排除二者)中,加入标点 skiplist 的模型在质量上略有胜出,并且让这套数据的 document 索引大小免费缩小了 9.6%:

import string

# model[2] 是 MultiVectorMask 模块
model[2].skiplist_words = list(string.punctuation)
model[2].resolve_with_tokenizer(model.tokenizer)  # token ids 会被缓存,因此修改后需要重新解析

基于基础 Transformer 构建模型

你也可以将任意基础 Transformer 传给 MultiVectorEncoder,系统会为模型追加一个全新、随机初始化的 token 级 projection:

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。投影层是随机初始化的,所以这个模型必须经过训练才能用。有意思的是,即便用很强的稠密向量 backbone,这套方法同样奏效。我在实验里给 Alibaba-NLP/gte-modernbert-base 重新接了一个投影层,仅凭这层投影加上 25k 训练对,结果就和现有 checkpoint 的起点只差了 0.03。

ColBERT 的那些经典 token 化技巧([MASK] 查询扩展、[Q] / [D] 前缀 token、文档长度上限、标点跳过列表)默认都是关闭的,可以自行配置。完整选项见 Creating Custom Models。就我个人经验来说,我在自己领域的微调里把 [MASK] 查询扩展试了四种配置,没有一种带来可观测的差异,所以不必执着于这套经典配方。

该选哪个起点?

为写这篇博客我专门测了一轮:挑了六个起点,用相同的训练配置各跑一遍,数据是 MIRIAD 里 25k 条医学问答-段落对,然后用 1,000 道留出的问题在 50,000 篇段落构成的语料上评估:

这个结果让我很意外,而且在两个模型系列上都得到了复现:*-unsupervised 版本的检查点在迁移到新领域时远胜于已训练完成的版本,即便起点更低也能反超。这些检查点处于大规模对比预训练之后、通用检索有监督微调之前的阶段,因此保留了所有的 late-interaction 结构,又不带通用检索微调带来的偏差——而后者恰恰是领域训练需要去抵消的东西。相比之下,已完成的检查点几乎不动,甚至出现回退,我试过的每一个学习率都是如此。

所以,如果你看中的模型系列发布了有监督训练前的检查点,就从它开始。如果没发布,那么在强检索预训练骨干网络上加一个新的投影头是次优选择。继续从完全训练完成的检查点出发,虽然最直觉,但却是领域适配效果最差的方案。

数据集

MultiVectorEncoderTrainer 使用 datasets.Datasetdatasets.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
})
"""

这就是本篇博客中要用的训练数据:来自 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")

如果本地数据需要预处理,可以用 datasets.Dataset.from_dict,通过一个由列表组成的字典来构建数据集:

from datasets import Dataset

queries = []
documents = []
# 打开文件,进行预处理、过滤、清洗等操作
# 然后追加到列表中

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

数据集格式

数据集格式必须与损失函数相匹配(或者反过来说,要选一个匹配你数据格式的损失函数)。验证某种数据集格式能否与某个损失函数配合,需要检查两步:

  1. 如果根据 Loss Overview 表格,你的损失函数需要 Label,那么数据集里就必须有一个 名为 "label" 或 "score" 的列。这一列会自动被当作标签。
  2. 根据 损失函数概览表,所有名称不是“label”或“score”的列都会被视为输入。剩余列的数量必须与所选损失函数的有效输入数量一致。这些列的名称无关紧要,只有顺序重要

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

  • 位置化 query 和 document 分配:第一列作为 query 嵌入,其余所有列作为 documents 嵌入,与列名无关。对于每列,都可以通过标准训练参数 router_mapping 覆盖这一默认设置。
  • 知识蒸馏格式:每个候选 document 占一列,即 (query, document_1, ..., document_N, scores),其中 scores 是每行对应的 N 个 teacher score 列表。对于将 query 和 document 的 ID 与文本数据集分开存储的 KD 数据集(例如 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,  # 每个分块编码的文档数量:限制内存,而非影响质量
)

mini_batch_size 参数通过按指定数量分块编码文档来限制内存,而有效对比学习批次大小可以自由选择;下面的实验中该值为 128,在我的消融实验里,继续增大批次并未带来额外收益。GradCache 保证无论分块大小如何,训练结果都完全一致,因此 GPU 内存较小时可以减小该值,代价只是运行时间增加。如果文档长度差异很大,可以考虑使用同系列的 mini_batch_num_tokens:它根据总 token 预算填充每个分块,而不是限制文档数量,因此包含异常超长文档的分块也不会造成内存峰值(我的实验中,mini_batch_size=16 且每个文档约 940 个 token,对应 mini_batch_num_tokens=15_000)。

多向量模型有一个特有的坑:对比损失的默认 scale=1.0,而稠密嵌入模型默认的是 scale=20.0。之所以用 20.0,是因为余弦相似度只是 [-1, 1] 区间内的一个值,范围太窄,撑不起尖锐的 softmax。而 MaxSim 分数是对每个查询 token 取一个最佳匹配的相似度再求和,因此天然就在 [0, query_length] 这个量级:32 个 token 的查询最高能打到 32。所以别直接把稠密训练脚本里的 scale=20.0 抄过来,那样会把 softmax 推到饱和,梯度直接消失。

如果想用更强的教师模型做蒸馏(目前最强的通用 late-interaction 模型基本都是这么训出来的),可以参考 MultiVectorDistillKLDivLoss,以及 Training Overview 文档里的 Knowledge Distillation 部分。

训练参数

你可以通过 MultiVectorEncoderTrainingArguments 类来定制训练过程。这个类允许你调整那些会影响训练速度、并帮助你理解训练过程的参数。

想了解最有用的训练参数,建议阅读 Multi-Vector Encoder > Training Overview > Training Arguments,值得花时间看看,能帮你把训练效果拉到最佳。

下面给个示例,用的是我实际训练时的参数:

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] "},  # checkpoint 自带的标记,按训练列名映射
    fp16=False,  # 如果你的 GPU 支持 FP16,可以设为 True
    bf16=True,   # 如果你的 GPU 支持 BF16,可以设为 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 等工具
)

其中几项参数值得说明一下:

  • prompts:训练过程不会自动套用模型里保存的 prompt,所以需要你手动把它映射到对应的训练列。这里把 checkpoint 自带的 [Q] 标记绑到问题列、[D] 绑到段落列,保证训练和推理时行为一致。
  • max_length(故意未设置):这个参数只会在训练阶段限制 token 化长度,适合你希望训练比模型实际服务长度更省钱的场景。我在这份数据上测过这个捷径的代价:把训练长度降到 512 tokens,速度大约能快 2 倍,但 NDCG@10 会掉约 0.015,而且这个差距不会随数据量增加而缩小,因为被截掉的内容模型根本看不到。除非你更在意速度而非质量,否则建议留空,让训练与推理保持一致。
  • learning_rate=1e-4:在 5e-6 到 2e-4 之间做了搜索后,这个偏高的学习率效果最好。

Evaluator(评估器)

如果想在训练过程中追踪模型表现,可以给 trainer 传入 eval_dataset 来观察评估损失,不过具体的检索指标会更有参考价值。Sentence Transformers 为多向量模型内置了以下评估器:

起点 Zero-shot NDCG@10 训练 25k 对之后 Delta
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
评估器 所需数据
MultiVectorInformationRetrievalEvaluator 查询、语料库及相关文档映射
MultiVectorNanoBEIREvaluator 无需数据
MultiVectorTripletEvaluator (锚例、正例、负例) 三元组
MultiVectorRerankingEvaluator {'query': '...', 'positive': [...], 'negative': [...]} 字典列表
MultiVectorDistillationEvaluator 带候选文档及教师模型分数的查询

对于领域微调来说,由你自己留出的数据构建的 MultiVectorInformationRetrievalEvaluator 才是关键。构建时有个小建议:语料库要有足够的区分度,才能把不同模型的差异拉开。以我的情况为例,MIRIAD 的问题是由其原文段落直接生成的,所以检索起来异常简单。在仅针对那 1 万条标准段落时,几乎所有模型的 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")

# 标准答案:1,000 个评估问题,每个问题对应一个段落,
# 评估集全部约 1 万个唯一段落作为初始语料库
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"]]}

# 干扰项:训练集中的唯一段落,用于让检索场景更真实
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)

Trainer

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

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 做 in-batch 负采样:有效 batch 大,但显存占用受 mini-batch 控制
    loss = CachedMultiVectorMultipleNegativesRankingLoss(model=model, mini_batch_size=16)

    # 6. 一个轻量的 dev 评估器,用于在训练过程中观察进度:500 条留出问题的评估,
    # 对应 eval 切分中约 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,  # 如果你的 GPU 支持 FP16,设为 True
        bf16=True,  # 如果你的 GPU 支持 BF16,设为 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 并开始训练
    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{traceback.format_exc()}")

if __name__ == "__main__":
    main()

整个配方就这么简单:一个预监督 checkpoint、一百万领域配对、批内负例、完整文档长度、再加一个偏高的学习率。我在单卡 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 的情况下,跟踪训练过程中的碳排放

通过 report_to 训练参数启用这些回调,例如 report_to=["wandb", "codecarbon"],并提前装好对应依赖即可。该参数默认为 "none",设为 report_to="all" 则会启用所有已安装依赖对应的集成。

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

多数据集训练

通常来说,性能顶尖的通用模型都是在多个数据集上同时训练得到的。不过由于各数据集的格式差异,这种做法并不容易实现。好在 MultiVectorEncoderTrainer 允许你在多个数据集上训练,并且不要求它们格式统一。此外,它还支持为不同数据集指定不同的损失函数。一次性训练多个数据集的步骤如下:
  • 使用一个 datasets.Dataset 实例的字典(或一个 datasets.DatasetDict)作为 train_dataset(可选地也作为 eval_dataset)。
  • (可选)使用一个损失函数的字典,将数据集名称映射到对应的损失函数。只有当你希望为不同数据集使用不同损失函数时才需要设置。
每个训练/评估批次只包含来自某一个数据集的样本。从多个数据集中采样批次的顺序由 MultiDatasetBatchSamplers 枚举定义,可以通过 multi_dataset_batch_sampler 参数传递给 MultiVectorEncoderTrainingArguments。可选值包括:
  • MultiDatasetBatchSamplers.ROUND_ROBIN:对每个数据集进行轮询采样,直到其中一个数据集耗尽。使用该策略时,每个数据集很可能无法用完所有样本,但各数据集的采样次数是均等的。
  • MultiDatasetBatchSamplers.PROPORTIONAL(默认):按照各数据集的大小比例进行采样。使用该策略时,每个数据集的所有样本都会被使用,且较大的数据集被采样得更频繁。

评估

为了确定微调模型的水平,我在 MIRIAD 评估集上对比了四种架构系列、共 50 多种检索模型配置。该评估集与上文 Evaluator 部分的构建方式完全相同:包含 1,000 个留出的医学问题,以及 200,000 个不同的段落(其中包括训练集里的 10,000 个金标准段落,并将其隐藏在 190,000 个去重干扰段落中)。这个语料库是如何选择起始模型?中 50,000 段语料库的 4 倍,因此两个表格中的分数无法直接比较。

NDCG@10 versus active parameters on the MIRIAD 200k benchmark, with an arrow marking the finetuning jump from mLateOn-unsupervised to mLateOn-medical

主要结果如下,完整表格可在下方展开查看:

模型 系列 NDCG@10
multi-vector-encoder/mLateOn-medical(我的模型) 多向量,微调 0.9139
lightonai/mLateOn 多向量,零样本 0.8520
lightonai/GTE-ModernColBERT-v1(解除上限) 多向量,零样本 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%,将首位错误率削减了超过三分之一。

架构规律同样一目了然:表格前列清一色是 late interaction 模型。在长文档场景下,即便训练数据和骨干网络完全相同,逐 token 的多向量表示也优于逐文档的单向量表示。DenseOn 和 LateOn 除了检索头以外,训练数据和架构完全一致,而 late interaction 版本领先 +0.12;多语言版本(mDenseOn 和 mLateOn)同样复现了这一差距,也是 +0.13。单向量模型靠规模也救不回来。Qwen3-Embedding-4B 是最强的稠密模型,其活跃参数(不含嵌入层)大约是我的模型的 33 倍,但仍然差了 0.13;而 8B 版本的表现反而不如 4B 版本。

BM25 的表现也出人意料地好,它击败了所有稀疏模型、所有因截断而受限的多向量模型,以及除三个稠密模型以外的所有模型:参数达数十亿的 Qwen3-Embedding-4B8B,以及 voyage-4-nano——后者读取完整的 32k token 上下文,也仅以 0.006 的微弱优势领先。不过别指望这套结果能直接复现到你的数据上。MIRIAD 的问题是从段落本身生成的,因此查询与目标段落之间的词面重叠远高于典型检索任务;同时 BM25 没有上下文长度限制,能利用每一个重叠词,而大多数神经模型都做了截断。BM25 基线成本极低、值得一跑,但不要迷信这里的领先幅度。

全部模型的概览,按分数排序,按架构族着色。

MIRIAD 200k 基准上所有评估模型的 NDCG@10 排序,按架构族着色

点击查看完整评估表
Model Family NDCG@10 acc@1
multi-vector-encoder/mLateOn-medical(本文模型) Multi-vector,finetuned 0.9139 0.849
lightonai/mLateOn Multi-vector 0.8520 0.758
lightonai/GTE-ModernColBERT-v1 @1024 Multi-vector 0.8502 0.763
lightonai/LateOn @1024 Multi-vector 0.8485 0.760
lightonai/mLateOn-unsupervised Multi-vector 0.8304 0.733
mixedbread-ai/mxbai-edge-colbert-v0-32m @1024 Multi-vector 0.8186 0.727
Qwen/Qwen3-Embedding-4B Dense 0.7817 0.669
Qwen/Qwen3-Embedding-8B Dense 0.7747 0.654
perplexity-ai/pplx-embed-v1-late-0.6b @1024 Multi-vector 0.7702 0.632
lightonai/ColBERT-Zero Multi-vector 0.7613 0.675
LiquidAI/LFM2.5-ColBERT-350M Multi-vector 0.7582 0.664
voyageai/voyage-4-nano Dense 0.7563 0.638
BM25 Lexical 0.7501 0.641
jinaai/jina-embeddings-v5-text-small-retrieval Dense 0.7470 0.620
Qwen/Qwen3-Embedding-0.6B Dense 0.7408 0.620
perplexity-ai/pplx-embed-v1-0.6b Dense 0.7384 0.615
mixedbread-ai/mxbai-edge-colbert-v0-32m Multi-vector 0.7350 0.639
mixedbread-ai/mxbai-edge-colbert-v0-17m Multi-vector 0.7271 0.631
answerdotai/answerai-colbert-small-v1 @512 Multi-vector 0.7264 0.615
lightonai/DenseOn @1024 Dense 0.7239 0.597
lightonai/mDenseOn @1024 Dense 0.7227 0.585
jinaai/jina-embeddings-v5-text-nano-retrieval Dense 0.7206 0.587
microsoft/harrier-oss-v1-0.6b Dense 0.7126 0.572
Alibaba-NLP/gte-modernbert-base Dense 0.7102 0.582
Snowflake/snowflake-arctic-embed-l-v2.0 Dense 0.7068 0.568
perplexity-ai/pplx-embed-v1-late-0.6b Multi-vector 0.7008 0.570
google/embeddinggemma-300m Dense 0.7000 0.563
lightonai/DenseOn Dense 0.6943 0.570
naver/splade-v3 Sparse 0.6853 0.574
ibm-granite/granite-embedding-small-english-r2 Dense 0.6813 0.546
naver/splade-v3-distilbert Sparse 0.6806 0.567
codefuse-ai/F2LLM-v2-0.6B Dense 0.6799 0.536
colbert-ir/colbertv2.0 @512 Multi-vector 0.6785 0.571
prithivida/Splade_PP_en_v1 Sparse 0.6755 0.577
lightonai/LateOn Multi-vector 0.6713 0.561
tomaarsen/embeddinggemma-300m-miriad-unsloth Dense, finetuned 0.6705 0.530
lightonai/LateOn-regularized Multi-vector 0.6673 0.554
lightonai/LateOn-unsupervised Multi-vector 0.6672 0.553
lightonai/GTE-ModernColBERT-v1 Multi-vector 0.6612 0.555
opensearch-project/opensearch-neural-sparse-encoding-v2-distill Sparse 0.6518 0.531
nomic-ai/nomic-embed-text-v1.5 (prompted) Dense 0.6387 0.498
mixedbread-ai/mxbai-embed-large-v1 Dense 0.6355 0.502
BAAI/bge-large-en-v1.5 Dense 0.6308 0.498
jinaai/jina-colbert-v2 @1024 多向量 0.6218 0.504
nomic-ai/nomic-embed-text-v1.5 密集 0.6203 0.487
answerdotai/answerai-colbert-small-v1 多向量 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
jinaai/jina-colbert-v2 多向量 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
colbert-ir/colbertv2.0 多向量 0.4347 0.346

标记为 @N 的模型在评测时把文档长度上限放宽到了 N 个 token,因为它们原本的上限(180 到 512 token)会截断平均长度达 941 token 的段落。对于每个多向量模型来说,这种放宽相比原始配置能带来 +0.08 到 +0.24 的 NDCG@10 提升,连单向量模型 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)

我在训练完成后的模型上做了后测,没有进行任何针对池化的训练,而在长文档上它的开销低得惊人。

20 万文档语料库的嵌入大小与 NDCG@10 对比,token pooling 的轨迹将多向量索引拉入了稠密模型的领地

实心点是未经压缩的嵌入,这样每个系列都能以相同方式计数并使用精确检索评分。不过实际部署不会直接用这些点:稠密索引通常配合 int8 或二值量化加 rescoring,稀疏索引会压缩倒排表,而多向量索引使用 PLAID 风格的残差压缩。读这些点时要把它当作相对存储成本,而不是你需要买的硬盘。

实线代表 token pooling。向量数减半只损失 0.0033 的 NDCG@10,rank-1 准确率不变;只保留四分之一的向量(11.2 GB)也能拿到 0.8991 的分数。曲线还会继续往下走(我测到了十分之一向量数,仍然有 0.8765),但量化方案一旦就位就没必要再往池化这个方向死磕了——这正是下面虚线要讨论的事。

虚线展示的是一个真实的部署方案。我提前把模型和基准交给了 Omar Khattab,他用 fast-plaid 测了下面这些配置:1-bit 残差量化、紧凑的 17-bit 中心点 id 和 18-bit 文档 id(替代通常未压缩的 64-bit 整数),再加上文档侧的剪枝:

配置 保留向量比例 索引大小 NDCG@10
1-bit PLAID,全部向量 100% 3.37 GB 0.8984
1-bit PLAID + 剪枝 65% 2.23 GB 0.8830
1-bit PLAID + 剪枝 42% 1.45 GB 0.8642

第一行的体积只有原始 embedding 的 1/13,但 NDCG@10 只掉到 0.0155。这比池化曲线上任何一点的取舍都划算得多。量化压缩每个向量的大小,池化和剪枝则削减向量数量,两者可以叠加,而量化应当是首选。继续往下,最后一行压缩到 1.45 GB,比 Qwen3-Embedding-8B 的 fp16 embedding(1.64 GB)还要小,同时得分高出 0.0895。"多向量索引太大"这个反对意见在配置得当的索引面前站不住脚。

这里的剪枝只是粗略实现的,目的是验证在量化之上做 token 削减是可行的,所以请把最后两行当作下限而非前沿。如果你完全不想手动调参,配套文章中索引构建一节涵盖了 fast-plaid、Qdrant、Weaviate 和 Vespa。

多向量检索的成本完全取决于索引。这个语料的原始 embedding 是 45 GB,而配置得当的索引体积至少能缩小 7 倍,精度却几乎不变。索引值得你投入与模型权重同等的关注。

致谢

感谢 Omar Khattab 测量了索引优化一节中量化与剪枝的索引配置,并就 late-interaction 索引的开销问题进行了讨论。

补充资源

训练示例

以下页面提供了带讲解的训练示例和训练脚本链接,可以帮助你熟悉多向量训练流程:

  • MIRIAD:面向医疗检索的领域训练,是本文方案的更早期、更简洁版本
  • MS MARCO:对比学习与知识蒸馏方案
  • 多模态:ColPali 风格的视觉文档检索训练
  • PEFT 适配器:基于 LoRA 的参数高效微调
原始来源: HuggingFace博客

评论 (0)