跳到正文
原文
Hugging Face Blog·· 2026-08-26精选AI 评分62

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

Training and Finetuning Multi-Vector Embedding Models with Sentence Transformers

AI 摘要

Sentence Transformers v6.0 新增 MultiVectorEncoder 模型类型,支持 ColBERT 式后期交互检索,并配套完整训练流程。

正文 · AI 翻译

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

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

最后,在评估部分,我将向你展示,我微调的multi-vector-encoder/mLateOn-medical模型——在单张 RTX 3090 上训练了 14.5 小时,与这篇博客文章同步完成——在我的医学检索评估中轻松超越了所有我能找到的通用检索模型:无论是稠密、稀疏、词汇还是多向量模型。

NDCG@10 on MIRIAD versus active parameters: the finetuned mLateOn-medical reaches the top at a fraction of the size of the strongest general-purpose models

如果你有兴趣微调稠密嵌入模型、稀疏嵌入模型或重排序模型,那么可以考虑阅读我之前关于训练和微调嵌入模型、训练和微调稀疏嵌入模型以及训练和微调重排序模型的博客文章。

这篇博客文章是关于训练多向量模型的。如果你想了解如何使用它们,从加载和编码到在向量数据库中建立索引,请参阅配套的使用 Sentence Transformers 的多向量(后期交互)嵌入模型博客文章。

目录

什么是多向量模型?

稠密嵌入模型将整个文本压缩成一个向量,相似度就是两个这样的摘要之间的一个点积。多向量模型(也称为后期交互或 ColBERT 风格模型)跳过了这种压缩。它为每个 token 保留一个小向量,并使用 MaxSim 算子对查询与文档进行评分,其中每个查询 token 找到其最佳匹配的文档 token,然后将分数相加。Token 级别的匹配恰好保留了单个向量不得不平均掉的细粒度信号,这通常意味着更强的检索能力,代价是更大的索引。

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

Dense embedding versus multi-vector late interaction

为什么要微调?

微调多向量模型能显著提升其在你的特定领域中的检索性能:词汇、查询风格以及相关性的定义在网页搜索、法律取证、代码搜索和科学文献综述之间都各不相同。由于查询和文档是逐 token 匹配的,多向量模型能够捕捉到单向量模型往往会平均掉的细粒度领域信号,而且即使只有少量领域内微调数据,它们也能有非常好的响应。

除此之外,大多数已发布的检索模型都是针对短段落配置的。经典的 ColBERT 检查点会将文档截断到 180 或 300 个 token,许多流行的稠密模型则截断到 256 或 512,因为它们基于 MS MARCO 风格的训练数据很少超过这个长度。如果你的文档很长,这些模型会在打分之前悄悄丢弃每篇文档的大部分内容。在我使用平均 941 个 token 的段落进行的医学评估中,我测得这种截断带来的损失高达 0.24 NDCG@10,远超过任何模型架构之间的差异。当你训练自己的模型时,你可以配置你的数据所需的文档长度。

LightOn 在代码检索中也遇到了同样的情形,通用的 LateOn 不够用,于是他们训练了 LateOn-Code。你的领域,无论是医学、法律、金融,还是你公司的内部文档,都不会有官方模型。这篇博文将向你展示如何自己构建它,只需几个小时,在一块消费级 GPU 上即可完成。

训练组件

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

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

让我们逐一仔细看看每个组件。

模型

多向量训练让你可以真正选择起点,而这一点比你想象的更重要。

微调现有的多向量模型

如果你想进一步微调现有的多向量模型,你完全不必操心架构:

from sentence_transformers import MultiVectorEncoder

# Loading in fp32 is preferred for training if your memory can handle it
model = MultiVectorEncoder(
    "lightonai/mLateOn-unsupervised",
    model_kwargs={"torch_dtype": "float32"},
    processor_kwargs={"model_max_length": 8192},  # the tokenizer-level token limit
)

检查点自带它的配方:它的查询和文档标记 token、它的投影头、它的打分跳过列表。对于微调,你通常希望保留所有这些,只改变你的数据所要求的部分。首先要检查的是长度配置,因为许多已发布的检查点将文档限制在 180 到 512 个 token(参见 为什么要微调?),而我的医学段落长达 1,400 个 token。mLateOn 系列已经支持骨干网络完整的 8192 token 上下文,但如果你起始的检查点带有上限,就把它们提高:

# Let the model read full documents instead of the caps it was trained with,
# e.g. GTE-ModernColBERT-v1 ships with query_length=48 and document_length=300
model[0].query_length = None
model[0].document_length = None

在取消各任务的上限后,截断会回退到分词器的 model_max_length,这就是为什么我在上面的加载时配置了该限制。

我又做了一处改动,添加了一个标点跳过列表,将标点符号 token 从文档侧评分和存储中排除。在四路消融实验(无、标点、停用词、两者)中,它在质量上略有胜出,并且在此数据上免费将文档索引缩小了 9.6%:

import string

# model[2] is the MultiVectorMask module
model[2].skiplist_words = list(string.punctuation)
model[2].resolve_with_tokenizer(model.tokenizer)  # token ids are cached, so re-resolve after changing

从基础 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 上全新投影仅凭投影和 25k 训练对,就达到了与现有检查点起点相差 0.03 以内的水平。

经典的 ColBERT 分词技巧([MASK] 查询扩展、[Q] / [D] 前缀 token、文档长度上限、标点跳过列表)默认全部关闭且可配置。完整列表请参见 Creating Custom Models。值得一提的是,我在我的领域微调中测试了四种配置下的 [MASK] 查询扩展,没有一种产生可测量的差异,所以不必觉得必须采用经典配方。

你应该选择哪个起点?

我在准备这篇博客文章时直接测量了这一点,选取六个起点,用相同的配方在来自 MIRIAD 的 25k 医学问题-段落对上分别训练,然后在 1,000 个留出问题上针对 50,000 段落语料库进行评估:

起点 零样本 NDCG@10 25k 对之后 差值
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 -

结果令我惊讶,并且在两个模型家族中都得到了复现。*-unsupervised 检查点对新领域的适应能力远好于它们已完成的同类,尽管起点更低却实现了超越。这些检查点位于大规模对比预训练之后、但在通用检索的监督微调之前,因此它们携带了所有后期交互结构,却没有那些领域训练随后必须撤销的通用调优。相比之下,已完成的检查点在我尝试的每个学习率下几乎纹丝不动,甚至出现退化。

所以,如果你喜欢的模型家族发布了监督前检查点,就从那里开始。如果没有,在强大的检索预训练骨干上全新投影是紧随其后的次优选择。从完全完成的检查点继续是领域适应最弱的选择,尽管它感觉上最自然。

数据集

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
})
"""

这是我在本博文中将用于训练的数据集:来自 MIRIAD 的 440 万个医学问题,每个问题都配有其答案所在的源段落(平均 941 个 token)。像这样简单的(查询,相关段落)对是为自己的领域收集的最容易的检索训练数据,而且正如你将看到的,它们就是你所需要的一切。

本地数据

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

from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_file.csv")
# or
dataset = load_dataset("json", data_files="my_file.json")

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

from datasets import Dataset

queries = []
documents = []
# Open a file, perform preprocessing, filtering, cleaning, etc.
# and append to the lists

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

数据集格式

重要的是,你的数据集格式要与你的损失函数匹配(或者你选择一个与数据集格式匹配的损失函数)。验证数据集格式是否适用于某个损失函数需要两个步骤:

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

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

  • 位置查询与文档分配:第一列被嵌入为 query,所有后续列被嵌入为 documents,无论列名如何。此默认设置可以通过标准的 router_mapping 训练参数按列覆盖。
  • 知识蒸馏格式:每个候选文档一列,即 (query, document_1, ..., document_N, scores),其中 scores 是每行 N 个教师分数的列表。对于将查询和文档 ID 与单独文本数据集一起存储的 KD 数据集(例如 lightonai/ms-marco-en-bge),你可以使用 resolve_ids 即时将 ID 解析为文本。

损失函数

损失函数量化模型在给定一批数据上的表现如何,使优化器能够更新模型权重以产生更有利(即更低)的损失值。适合你任务的损失函数取决于你拥有的数据以及你想要实现的目标。你可以在 Loss Overview 中找到完整的选项列表。

对于问答或问题-段落对的常见情况,主力方法是使用 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,  # how many documents to encode per chunk: bounds memory, not quality
)

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 饱和并杀死你的梯度。

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

训练参数

你可以使用 MultiVectorEncoderTrainingArguments 类来自定义训练过程。该类让你调整可能影响训练速度的参数,并帮助你理解训练过程中发生了什么。

关于最有用的训练参数的更多信息,请查看 多向量编码器 > 训练概览 > 训练参数。值得一读,以充分利用你的训练。

下面是一个示例,使用了我实际训练运行中的值:

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

args = MultiVectorEncoderTrainingArguments(
    # Required parameter:
    output_dir="models/mLateOn-medical",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=128,  # the effective contrastive batch, thanks to GradCache
    per_device_eval_batch_size=16,
    learning_rate=1e-4,
    warmup_steps=0.05,
    prompts={"question": "[Q] ", "passage_text": "[D] "},  # the checkpoint's markers, keyed by training column
    fp16=False,  # Set to True if you have a GPU that supports FP16
    bf16=True,  # Set to True if you have a GPU that supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # in-batch negatives benefit from no duplicates
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=0.1,
    save_strategy="steps",
    save_steps=0.05,
    logging_steps=0.01,
    run_name="mLateOn-medical",  # Will be used in e.g. Trackio, W&B, etc.
)

其中几项值得评论:

  • prompts:训练不会自动应用模型中存储的提示,因此请将它们显式映射到你的训练列上。这里就是检查点的 [Q] 标记用于问题列,[D] 用于段落列,使训练与推理保持一致。
  • max_length(故意不设置):此参数仅在训练期间限制分词,适用于你希望训练比模型完整服务长度更便宜的情况。我测量了这个捷径在此数据上的代价。以 512 token 训练损失了约 0.015 NDCG@10,换来约 2 倍的速度,而且这一差距并未随更多数据而缩小,因为模型根本看不到被截断的部分。除非你需要速度胜过质量,否则保持不设置,使训练与推理匹配。
  • 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 evaluation questions, each mapping to its own passage, with the
# eval split's full ~10k unique passages as the initial corpus
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"]]}

# Distractors: unique train passages that make the haystack realistic
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. Load the starting checkpoint: contrastively pretrained, not yet supervised
    # Loading in fp32 is preferred for training if your memory can handle it
    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. Lift the per-task length caps so training and inference see full medical passages
    model[0].query_length = None
    model[0].document_length = None

    # 3. Skip punctuation tokens during scoring: a small quality win and a 9.6% smaller index
    model[2].skiplist_words = list(string.punctuation)
    model[2].resolve_with_tokenizer(model.tokenizer)

    # 4. Load 1 million medical question-passage pairs
    train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train").select(range(1_000_000))

    # 5. In-batch negatives with GradCache: large effective batch, memory-bounded chunks
    loss = CachedMultiVectorMultipleNegativesRankingLoss(model=model, mini_batch_size=16)

    # 6. A light dev evaluator to watch progress during training: 500 held-out questions
    # against the eval split's ~10k unique passages. The full 200k protocol runs afterwards.
    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. Training arguments, as discussed above
    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,  # Set to True if you have a GPU that supports FP16
        bf16=True,  # Set to True if you have a GPU that supports BF16
        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. Create a trainer & train
    trainer = MultiVectorEncoderTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        loss=loss,
        evaluator=dev_evaluator,
    )
    trainer.train()

    # 9. Save the trained model
    model.save_pretrained(f"models/{run_name}/final")

    # 10. (Optional) Push it to the Hugging Face Hub
    try:
        model.push_to_hub(run_name)
    except Exception:
        logging.error(f"Error uploading model to the Hugging Face Hub:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

这就是完整的配方:一个预监督检查点、一百万对领域数据、批内负样本、完整文档长度,以及比通常更高的学习率。这次运行在我的单张 RTX 3090 上耗时 14.5 小时,峰值显存 17.5 GB,而每一个选择都是经过测量比较后的胜出者,而非猜测。

对于预算较小的读者,我的缩放实验表明,10 万对数据(训练 75 分钟)与完整的百万对运行相比,NDCG@10 差距在 0.012 以内。大部分收益来自第一个小时。

Callbacks

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

通过 report_to 训练参数启用这些功能,例如 report_to=["wandb", "codecarbon"],并安装所需的依赖项。它默认为 "none",而 report_to="all" 会激活所有已安装依赖项的集成。

有关这些回调以及如何创建自己的回调的更多信息,请参阅 Transformers Callbacks 文档。

Multi-Dataset Training

通常,表现最佳的通用模型会同时在多个数据集上进行训练。然而,由于每个数据集的格式不同,这种方法可能具有挑战性。幸运的是,MultiVectorEncoderTrainer 允许你在多个数据集上训练,而无需统一格式。此外,它还提供了对每个数据集应用不同损失函数的灵活性。以下是同时使用多个数据集进行训练的步骤:

  • 使用 datasets.Dataset 实例的字典(或 datasets.DatasetDict)作为 train_dataset(也可选作为 eval_dataset)。
  • (可选)使用损失函数字典,将数据集名称映射到损失函数。仅当你希望为不同数据集使用不同损失函数时才需要。

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

  • MultiDatasetBatchSamplers.ROUND_ROBIN:从每个数据集轮流采样,直到一个数据集耗尽。使用此策略,可能不会使用每个数据集的所有样本,但每个数据集被采样的机会均等。
  • MultiDatasetBatchSamplers.PROPORTIONAL(默认):按每个数据集的大小比例进行采样。使用此策略,会使用每个数据集的所有样本,且较大的数据集被更频繁地采样。

Evaluation

为了了解微调模型的表现,我在 MIRIAD 评估集上,针对四个架构家族的 50 多种检索模型配置进行了评估,评估集的构建方式与上文 Evaluator 部分完全一致,使用 1,000 个留出的医学问题检索 200,000 个唯一段落(10k 个金标段落隐藏在来自训练集的 190k 个去重干扰项中)。该语料库是 Which starting point should you pick? 中 50,000 段落语料库的四倍,因此两个表格中的分数不可直接比较。

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

微调模型位居榜首,比任何架构中最强的零样本模型高出 +0.062 NDCG@10。换句话说,最强的零样本模型在 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 基线成本低廉,始终值得运行,只是不要指望这个优势。

完整阵容一览,按分数排序并按架构家族着色。

Sorted NDCG@10 on the MIRIAD 200k benchmark for every evaluated model, colored by architecture family

点击查看完整评估表

模型 家族 NDCG@10 acc@1
multi-vector-encoder/mLateOn-medical(我的) 多向量,微调 0.9139 0.849
lightonai/mLateOn 多向量 0.8520 0.758
lightonai/GTE-ModernColBERT-v1 @1024 多向量 0.8502 0.763
lightonai/LateOn @1024 多向量 0.8485 0.760
lightonai/mLateOn-unsupervised 多向量 0.8304 0.733
mixedbread-ai/mxbai-edge-colbert-v0-32m @1024 多向量 0.8186 0.727
Qwen/Qwen3-Embedding-4B 稠密 0.7817 0.669
Qwen/Qwen3-Embedding-8B 稠密 0.7747 0.654
perplexity-ai/pplx-embed-v1-late-0.6b @1024 多向量 0.7702 0.632
lightonai/ColBERT-Zero 多向量 0.7613 0.675
LiquidAI/LFM2.5-ColBERT-350M 多向量 0.7582 0.664
voyageai/voyage-4-nano 稠密 0.7563 0.638
BM25 词法 0.7501 0.641
jinaai/jina-embeddings-v5-text-small-retrieval 稠密 0.7470 0.620
Qwen/Qwen3-Embedding-0.6B 稠密 0.7408 0.620
perplexity-ai/pplx-embed-v1-0.6b 稠密 0.7384 0.615
mixedbread-ai/mxbai-edge-colbert-v0-32m 多向量 0.7350 0.639
mixedbread-ai/mxbai-edge-colbert-v0-17m 多向量 0.7271 0.631
answerdotai/answerai-colbert-small-v1 @512 多向量 0.7264 0.615
lightonai/DenseOn @1024 稠密 0.7239 0.597
lightonai/mDenseOn @1024 稠密 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
perplexity-ai/pplx-embed-v1-late-0.6b 多向量 0.7008 0.570
google/embeddinggemma-300m 稠密 0.7000 0.563
lightonai/DenseOn 稠密 0.6943 0.570
naver/splade-v3 稀疏 0.6853 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 @512 多向量 0.6785 0.571
prithivida/Splade_PP_en_v1 稀疏 0.6755 0.577
lightonai/LateOn 多向量 0.6713 0.561
tomaarsen/embeddinggemma-300m-miriad-unsloth 稠密,微调 0.6705 0.530
lightonai/LateOn-regularized 多向量 0.6673 0.554
lightonai/LateOn-unsupervised 多向量 0.6672 0.553
lightonai/GTE-ModernColBERT-v1 多向量 0.6612 0.555
opensearch-project/opensearch-neural-sparse-encoding-v2-distill 稀疏 0.6518 0.531
nomic-ai/nomic-embed-text-v1.5(带提示) 稠密 0.6387 0.498
mixedbread-ai/mxbai-embed-large-v1 稠密 0.6355 0.502
BAAI/bge-large-en-v1.5 稠密 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)

我是在训练完成的模型上事后测量的,没有进行任何池化感知训练,而在长文档上它的成本低得惊人。

Embedding size for the 200,000-passage corpus versus NDCG@10, with the token pooling trajectory sweeping the multi-vector index into dense-model territory

实心点表示未压缩的嵌入,因此每个系列都以相同方式计数,并用精确搜索评分。不过,你不会以那种方式部署它们中的任何一个。稠密索引通常使用 int8 或二值量化并重新评分,稀疏索引压缩其倒排列表,而多向量索引使用 PLAID 风格的残差压缩。不要把那些点理解为你需要购买的磁盘,而应理解为相对存储成本。

Token 池化是实线。将向量数量减半会损失 0.0033 NDCG@10,且 rank-1 准确率保持不变,而只保留四分之一,即 11.2 GB,仍能得分 0.8991。曲线还在继续(我测量到向量的十分之一,仍为 0.8765),但一旦量化提上日程,就没有什么理由把池化推到那么远,这正是下面虚线所涉及的内容。

虚线是真实部署可能的样子。我让 Omar Khattab 提前使用了该模型和基准测试,他用 fast-plaid 在 1 位残差量化下测量了这些配置,使用紧凑的 17 位质心 id 和 18 位文档 id,而不是其通常的未打包 64 位整数,外加文档侧剪枝:

配置 保留向量数 索引 NDCG@10
1 位 PLAID,全部向量 100% 3.37 GB 0.8984
1 位 PLAID + 剪枝 65% 2.23 GB 0.8830
1 位 PLAID + 剪枝 42% 1.45 GB 0.8642

第一行比原始嵌入小 13 倍,NDCG@10 为 0.0155。这比池化曲线上的任何位置都是更好的权衡。量化缩小每个向量,而池化和剪枝减少你保留的数量,因此它们可以叠加,而量化是首先应该采用的。进一步推进,最后一行落在 1.45 GB,小于 Qwen3-Embedding-8B 的 fp16 嵌入(1.64 GB),同时得分高出 0.0895。多向量索引太大的反对意见在正确配置的索引面前站不住脚。

这里的剪枝是朴素的,只是为了确立 token 缩减在量化之上仍然有效,因此请把底部两行视为下限而非前沿。如果你根本不想手动调优量化,配套文章的 Indexing 部分涵盖了 fast-plaid、Qdrant、Weaviate 和 Vespa。

多向量检索的成本只取决于其索引。该语料库的原始嵌入为 45 GB,而正确配置的索引在几乎相同准确率下至少小 7 倍。索引值得你像关注检查点一样关注。

致谢

感谢 Omar Khattab 在 优化索引 中测量了量化和剪枝的索引配置,并围绕后期交互索引成本进行了讨论。

其他资源

训练示例

这些页面包含带有解释的训练示例以及训练脚本的链接。你可以用它们来熟悉多向量训练循环:

  • MIRIAD:针对医学检索的领域特定训练,是本博文配方的更早、更简单的近亲
  • MS MARCO:对比学习和知识蒸馏配方
  • 多模态:ColPali 风格的视觉文档检索训练
  • PEFT 适配器:使用 LoRA 进行参数高效微调

文档

如需进一步学习,你也可以探索以下关于 Sentence Transformers 的资源:

这里还有一个你可能会感兴趣的高级页面:

以及配套博文,涵盖关于使用这些模型的所有内容:

来源:Hugging Face Blog · huggingface.co

00