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

AI资讯分类2个月前发布
546 00
广告也精彩

注意:许多可以直接与 Sentence Transformers 配合使用的公开数据集都已在 Hugging Face Hub 上被打上了 sentence-transformers 标签,因此你可以轻松地在 https://huggingface.co/datasets?other=sentence-transformers 上找到它们。不妨浏览一下这些数据集,寻找可能对你的任务、领域或语言有用的现成数据。

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

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

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,
})
模型家族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

检查点自带其专属配置:它的查询和文档标记 token、它的投影头、它的评分跳表。对于微调,你通常希望保留所有这些,只根据你的数据需求做改动。首先要检查的是长度配置,因为许多已发布的检查点将文档上限设为 180 到 512 个 token(参见“为什么要微调?”),而我的医学文本段落长达 1,400 个 token。mLateOn 系列已经支持骨干模型完整的 8192 token 上下文,但如果你的起始检查点带有长度上限,请解除它们:

  1. 如果你的损失函数根据“损失概览”表要求一个标签(Label),那么你的数据集必须有一个名为“label”或“score”的列。该列会被自动视为标签。
  2. 根据“损失概览”表,所有不叫“label”或“score”的列都被视为输入(Inputs)。剩余列的数量必须与你所选损失函数的有效输入数量一致。这些列的名称无关紧要,只有顺序才重要。
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
pooling = HierarchicalTokenPooling(pool_factor=4)
document_embeddings = model.encode_document(passages, token_pooling=pooling)

MultiVectorEncoderTrainer 支持各种 transformers.TrainerCallback 子类,包括:

BM25 的表现也出人意料地好,击败了所有稀疏模型、所有受截断限制的多向量模型,以及除三个稠密模型之外的所有模型:这三个例外是数十亿参数规模的 Qwen3-Embedding-4B 和 8B,以及 voyage-4-nano——后者通过读取完整的 32k token 上下文,以仅 0.006 的微弱优势险胜。不过,别指望这一优势能迁移到你自己的数据上。MIRIAD 的问题是从段落中生成的,因此查询与其黄金段落之间的词汇重叠远大于典型检索场景,而 BM25 不受上下文长度限制,可以利用所有这些重叠词,而大多数神经检查点则会截断。BM25 基线成本低廉,始终值得运行,只是别指望能有这么大的优势。

  • `MultiDatasetBatchSamplers.ROUND_ROBIN`:从每个数据集轮询采样,直到其中一个数据集被耗尽。使用此策略时,可能不会用到每个数据集中的所有样本,但每个数据集被采样的机会是均等的。
  • `MultiDatasetBatchSamplers.PROPORTIONAL`(默认):按各数据集大小比例进行采样。使用此策略时,每个数据集中的所有样本都会被使用,且较大的数据集会被更频繁地采样。
  • 提示词:训练时不会自动套用模型中存储的提示词,因此需要显式地将它们映射到你的训练列上。这里,检查点中的 [Q] 标记对应问题列,[D] 标记对应段落列,从而保持训练与推理的一致性。
  • max_length(刻意不设置):此参数仅在训练期间限制分词长度,适用于当你希望训练成本低于模型完整服务长度时的场景。我测量了这种简化方式在此数据上的代价。在 512 token 长度下训练,NDCG@10 大约损失 0.015,但速度提升约 2 倍,而且增加数据量也无法缩小这一差距,因为模型根本看不到被截断的内容。除非你更需要速度而非质量,否则请保持不设置,让训练与推理保持一致。
  • learning_rate=1e-4:在从 5e-6 到 2e-4 的范围内进行扫描后,我发现这个略高于常规的学习率效果最佳。

训练参数

起点零样本 NDCG@1025k 样本对训练后差值
lightonai/mLateOn-unsupervised0.90870.9398+0.0311
lightonai/mLateOn0.92770.9319+0.0042
lightonai/LateOn-unsupervised0.90260.9206+0.0180
lightonai/LateOn0.91850.9105-0.0080
lightonai/GTE-ModernColBERT-v10.91980.9007-0.0191
gte-modernbert-base 上的全新头部–0.9177–

多向量特有的一个陷阱是,对比损失的默认 `scale=1.0`,而密集嵌入的对应版本默认 `scale=20.0`。之所以存在 20.0,是因为余弦相似度是 [-1, 1] 中的单个值,对于尖锐的 softmax 来说范围太窄。而 MaxSim 分数则是对每个查询 token 求和一个最佳匹配相似度,因此其范围已经大致覆盖 [0, query_length]:一个 32 个 token 的查询最高可得 32 分。所以不要从密集训练脚本中照搬 `scale=20.0`,因为它会使 softmax 饱和并扼杀梯度。

如需进一步学习,你可能还想浏览以下关于 Sentence Transformers 的资源:

Token 池化是实线。向量数量减半仅损失 0.0033 的 NDCG@10,且排名第一的准确率完全不受影响;即使只保留四分之一的向量(11.2 GB),得分仍达 0.8991。曲线还在继续延伸(我测到只保留十分之一向量时,得分仍有 0.8765),但一旦引入量化,就没有太大必要把池化推到那么极致——这正是下方虚线所展示的内容。

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

我还做了一项改动,添加了一个标点跳表,将标点 token 从文档侧的评分和存储中排除。在 4 路消融实验(无排除、排除标点、排除停用词、两者都排除)中,它在质量上略有胜出,并且在此数据上将文档索引缩小了 9.6%,且无需额外成本:

致谢

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

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

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

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

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()

评估

微调现有的多向量模型

# 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

第一行比原始嵌入向量小 13 倍,NDCG@10 仅损失 0.0155。这个性价比远优于池化曲线上的任何一点。量化压缩每个向量的大小,而池化和剪枝减少保留的向量数量,因此三者可以叠加使用,而量化是最值得优先采用的手段。继续推进后,最后一行达到 1.45 GB,比 Qwen3-Embedding-8B 的 fp16 嵌入向量(1.64 GB)还小,同时得分高出 0.0895。所谓“多向量索引太大”的质疑,在正确配置的索引面前站不住脚。

这就是全部配方:一个预监督检查点、一百万条领域配对数据、批内负样本、完整文档长度,以及高于常规的学习率。该训练在我单张 RTX 3090 上耗时 14.5 小时,峰值显存占用 17.5 GB,而其中每一个选择都是经过实测对比胜出的结果,而非凭空猜测。

评估器所需数据
MultiVectorInformationRetrievalEvaluator查询、语料库及相关文档映射
MultiVectorNanoBEIREvaluator无需数据
MultiVectorTripletEvaluator(锚点、正例、负例)三元组
MultiVectorRerankingEvaluator{‘query’: ‘…’, ‘positive’: […], ‘negative’: […]} 字典列表
MultiVectorDistillationEvaluator带有候选文档和教师分数的查询
配置保留向量比例索引大小NDCG@10
1 位 PLAID,全部向量100%3.37 GB0.8984
1 位 PLAID + 剪枝65%2.23 GB0.8830
1 位 PLAID + 剪枝42%1.45 GB0.8642

对于预算有限的读者,我的扩展实验表明,10 万条配对数据(75 分钟训练时长)与完整百万配对数据运行的 NDCG@10 差距仅在 0.012 以内。大部分收益在第一个小时内即可获得。

为什么要微调?

文档

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

从基础 Transformer 构建一个

为了了解微调后模型的表现,我在 MIRIAD 评估集上将其与四个架构家族的 50 多种检索模型配置进行了对比评估。该评估集的构建方式与上文“评估器”部分完全相同,包含 1,000 个留出的医学问题,用于检索 200,000 个唯一段落(其中 10k 个黄金段落隐藏在训练集中去重后的 190k 个干扰项中)。该语料库规模是《你应该选择哪个起点?》一文中 50,000 段落语料库的四倍,因此两张表格中的分数不可直接比较。

微调后的模型位居榜首,比任何架构中最强的零样本模型高出 +0.062 NDCG@10。换句话说,最强的零样本模型在 75.8% 的查询中将正确段落排在第一位,而微调后的模型在 84.9% 的查询中做到了这一点,将排名第一的错误率降低了超过三分之一。

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

对多向量检索的合理质疑在于索引大小,而这个领域几乎是最糟糕的情况。我的模型为每个 token 存储一个向量,每个段落大约需要 878 个向量,因此 20 万段落的语料库在 fp16 下大约占用 45 GB,而稠密模型则远低于 1 GB。文档长度是造成如此巨大差距的原因。配套文章中的 Natural Questions 段落平均每个约 125 个 token 向量,少了七倍,因此短段落语料库的索引起点远小于本文的情况。HierarchicalTokenPooling 模块正是通过聚类每个文档的 token 嵌入并存储聚类均值来压缩这一点,大约保留 1 / pool_factor 的向量:

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

数据集

Image 4: Dense embedding versus multi-vector late interaction

不要低估在你的领域微调多向量模型的力量。在单块消费级 GPU 上花费十四个半小时,就产出了一个在该数据上任何通用检索器都无法企及的模型,而且配方只是一个脚本,没有教师模型,也没有挖掘负样本!

除此之外,还有两个多向量(multi-vector)特有的约定:

虚线代表真实部署中可能达到的效果。我提前向 Omar Khattab 开放了模型和基准测试的访问权限,他用 fast-plaid 在 1 位残差量化下测了这些配置,使用紧凑的 17 位质心 ID 和 18 位文档 ID(而非通常未压缩的 64 位整数),并配合文档侧剪枝:

本地数据

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
)

标记为 @N 的模型在评估时,其文档长度上限被提升至 N 个 token,因为其原生上限(180 到 512 个 token)会截断平均 941 个 token 的段落。对于每个多向量模型,这一提升相比按原样部署的行,带来了 +0.08 到 +0.24 的 NDCG@10 增益,即使是稠密模型 DenseOn 也因同样的处理获得了 +0.03 的提升。

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

为了在训练期间跟踪模型性能,你可以向训练器传入 eval_dataset 以获取评估损失,但具体的检索指标更具参考价值。Sentence Transformers 为多向量模型内置了以下评估器:

我在训练完成的模型上事后进行了测量,没有进行任何池化感知训练,在长文档上它非常节省资源。

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

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

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

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

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

目录

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

模型家族NDCG@10acc@1
multi-vector-encoder/mLateOn-medical(我的模型)多向量,微调0.91390.849
lightonai/mLateOn多向量0.85200.758
lightonai/GTE-ModernColBERT-v1 @1024多向量0.85020.763
lightonai/LateOn @1024多向量0.84850.760
lightonai/mLateOn-unsupervised多向量0.83040.733
mixedbread-ai/mxbai-edge-colbert-v0-32m @1024多向量0.81860.727
Qwen/Qwen3-Embedding-4B稠密0.78170.669
Qwen/Qwen3-Embedding-8B稠密0.77470.654
perplexity-ai/pplx-embed-v1-late-0.6b @1024多向量0.77020.632
lightonai/ColBERT-Zero多向量0.76130.675
LiquidAI/LFM2.5-ColBERT-350M多向量0.75820.664
voyageai/voyage-4-nano稠密0.75630.638
BM25词法0.75010.641
jinaai/jina-embeddings-v5-text-small-retrieval稠密0.74700.620
Qwen/Qwen3-Embedding-0.6B稠密0.74080.620
perplexity-ai/pplx-embed-v1-0.6b稠密0.73840.615
mixedbread-ai/mxbai-edge-colbert-v0-32m多向量0.73500.639
mixedbread-ai/mxbai-edge-colbert-v0-17m多向量0.72710.631
answerdotai/answerai-colbert-small-v1 @512多向量0.72640.615
lightonai/DenseOn @1024稠密0.72390.597
lightonai/mDenseOn @1024稠密0.72270.585
jinaai/jina-embeddings-v5-text-nano-retrieval稠密0.72060.587
microsoft/harrier-oss-v1-0.6b稠密0.71260.572
Alibaba-NLP/gte-modernbert-base稠密0.71020.582
Snowflake/snowflake-arctic-embed-l-v2.0稠密0.70680.568
perplexity-ai/pplx-embed-v1-late-0.6b多向量0.70080.570
google/embeddinggemma-300m稠密0.70000.563
lightonai/DenseOn稠密0.69430.570
naver/splade-v3稀疏0.68530.574
ibm-granite/granite-embedding-small-english-r2稠密0.68130.546
naver/splade-v3-distilbert稀疏0.68060.567
codefuse-ai/F2LLM-v2-0.6B稠密0.67990.536
colbert-ir/colbertv2.0 @512多向量0.67850.571
prithivida/Splade_PP_en_v1稀疏0.67550.577
lightonai/LateOn多向量0.67130.561
tomaarsen/embeddinggemma-300m-miriad-unsloth稠密,微调0.67050.530
lightonai/LateOn-regularized多向量0.66730.554
lightonai/LateOn-unsupervised多向量0.66720.553
lightonai/GTE-ModernColBERT-v1多向量0.66120.555
opensearch-project/opensearch-neural-sparse-encoding-v2-distill稀疏0.65180.531
nomic-ai/nomic-embed-text-v1.5 (prompted)稠密0.63870.498
mixedbread-ai/mxbai-embed-large-v1稠密0.63550.502
BAAI/bge-large-en-v1.5稠密0.63080.498
jinaai/jina-colbert-v2 @1024多向量0.62180.504
nomic-ai/nomic-embed-text-v1.5稠密0.62030.487
answerdotai/answerai-colbert-small-v1多向量0.61840.514
tomaarsen/splade-modernbert-base-miriad稀疏,微调0.61420.473
NeuML/biomedbert-base-colbert多向量0.59630.463
BAAI/bge-base-en-v1.5稠密0.59300.454
BAAI/bge-small-en-v1.5稠密0.58810.457
sentence-transformers/all-mpnet-base-v2稠密0.51590.396
jinaai/jina-colbert-v2多向量0.49920.401
mixedbread-ai/mxbai-colbert-large-v1多向量0.46900.358
sentence-transformers/static-retrieval-mrl-en-v1稠密0.46140.323
sentence-transformers/all-MiniLM-L6-v2稠密0.44580.321
colbert-ir/colbertv2.0多向量0.43470.346

训练示例

什么是多向量模型?

对于领域微调,基于你自己的留出数据构建的 MultiVectorInformationRetrievalEvaluator 才是关键。构建时的一个技巧是,语料库的难度要足够高,以便区分不同模型。以我的情况为例,MIRIAD 问题是从其对应的源段落生成的,这使得检索异常简单。仅针对 10k 个黄金段落,几乎所有模型的 NDCG@10 得分都高于 0.97。如果你的评估也出现这种饱和现象,可以加入干扰段落(我使用训练集中去重后的段落),直到分数分布拉开差距:

评估器

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({...})
# )

多向量训练让你在起点选择上拥有真正的自由度,而这一点的重要性可能超乎你的预期。

Image 7: 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

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

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

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

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

多数据集训练

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

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.
)

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

优化索引

损失函数

感谢 Omar Khattab 在“优化索引”部分测算了量化和剪枝后的索引配置,并参与了关于后期交互索引成本的讨论。

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

这里的剪枝是朴素实现,仅用于证明 token 缩减在量化之上依然有效,因此建议把下面两行视为下限而非最优水平。如果你完全不想手动调优量化参数,配套文章中的 Indexing 部分介绍了 fast-plaid、Qdrant、Weaviate 和 Vespa。

你应该选择哪个起点?

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

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

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

Sentence Transformers 是一个 Python 库,用于使用和训练嵌入向量模型与重排序模型,适用于检索增强生成、语义搜索、语义文本相似度等多种应用场景。其 v6.0 更新引入第四种模型类型:MultiVectorEncoder,用于 ColBERT 风格的后期交互检索,并配套提供了完整的训练方法。在本文中,我将展示如何使用它来微调一个在你的数据上优于通用检索模型的多向量模型。该方法还可以从头训练出强大的新多向量模型。以下所有内容均基于 `pip install -U “sentence-transformers[train]”` 运行。微调多向量模型涉及多个组件:模型本身、数据集、损失函数、训练参数、评估器和训练器类。我将逐一介绍这些组件,并附上如何使用它们微调出强大多向量模型的实用示例。

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

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

训练组件

更多资源

  • 分布式训练

Hugging Face Hub 上的数据

你可以使用 load_dataset 函数从 Hub 上的数据集中加载数据:

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

数据集格式

通过 report_to 训练参数启用这些回调,例如 report_to=[“wandb”, “codecarbon”],并确保已安装相应依赖。默认值为 “none”,而 report_to=”all” 会激活所有已安装依赖对应的集成。

  • MIRIAD:面向医学检索的领域专用训练,是本文配方的一个更早期、更简单的同类方案
  • MS MARCO:对比学习与知识蒸馏配方
  • 多模态:ColPali 风格的视觉文档检索训练
  • PEFT 适配器:使用 LoRA 进行参数高效微调

回调

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

  • 安装
  • 快速上手
  • 用法
  • 创建自定义模型
  • 预训练模型
  • 训练概览(本文是训练概览文档的精炼版)
  • 损失函数概览
  • API 参考

取消每个任务的长度上限后,截断会回退到 tokenizer 的 model_max_length,这就是为什么我在上面加载时配置了该限制。

如果你对微调稠密嵌入向量模型、稀疏嵌入向量模型或重排序模型感兴趣,建议阅读我之前撰写的《训练与微调嵌入向量模型》《训练与微调稀疏嵌入向量模型》以及《训练与微调重排序模型》等博客文章。

`mini_batch_size` 参数通过按此大小对文档进行分块编码来限制内存占用,而有效的对比批次大小(在我下面的运行中为 128,且在我的消融实验中更大的批次并未带来额外收益)仍可自由选择。GradCache 保证无论块大小如何,结果都完全一致,因此在较小的 GPU 上可以降低该值,代价仅仅是墙钟时间增加。当你的文档长度差异较大时,可以考虑其姊妹参数 `mini_batch_num_tokens`,它按总 token 预算而非文档数量来打包每个块,因此包含异常长文档的块绝不会导致内存激增(我的 `mini_batch_size=16`,每篇文档约 940 个 token,对应 `mini_batch_num_tokens=15_000`)。

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

点击查看完整评估表

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

多向量检索的成本完全取决于索引。该语料库的原始嵌入向量为 45 GB,而正确配置的索引在几乎相同准确率下至少小 7 倍。索引值得你投入与模型权重同等的关注。

模型

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
)
  • 使用 Sentence Transformers 构建多向量(后期交互)嵌入模型
  • 什么是多向量模型?

  • 为什么要微调?

  • 训练组件

  • 模型

    • 微调现有的多向量模型
    • 从基础 Transformer 构建多向量模型
    • 你应该选择哪个起点?
  • 数据集

    • Hugging Face Hub 上的数据
    • 本地数据
    • 数据集格式
  • 损失函数

  • 训练参数

  • 评估器

  • 训练器

    • 回调函数
    • 多数据集训练
  • 评估

    • 优化索引
  • 致谢

  • 更多资源

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

对于常见的问答对或问题-段落对场景,最常用的方法是使用 MultiVectorMultipleNegativesRankingLoss 进行批内负样本训练,其中批次中的每个其他文档都充当每个查询的负样本。更大的批次意味着更多的负样本和更强的训练效果,因此在实际操作中你会希望使用其 GradCache 变体,即 CachedMultiVectorMultipleNegativesRankingLoss,它可以将有效批次大小与 GPU 显存容量解耦。

这是我在本博客文章中用来训练的数据集:来自 MIRIAD 的 440 万个医学问题,每个问题都配对了包含其答案的源段落(平均 941 个 token)。像这样的简单(查询,相关段落)配对,是为你自己的领域收集检索训练数据最容易的方式,而且你会看到,它们已经足够了。

主要结果如下,完整表格见下方折叠区域:

你也可以将 MultiVectorEncoder 指向任何基础 Transformer,系统会自动为你追加一个全新的、随机初始化的 token 级投影层:

架构模式同样一目了然,表格顶部的模型清一色采用晚期交互(late interaction)。在长文档上,每个 token 一个向量要胜过每篇文档一个向量,即便在训练数据和骨干网络都匹配的情况下也是如此。DenseOn 和 LateOn 共享训练数据和架构,仅头部不同,而晚期交互的孪生模型以 +0.12 的优势胜出;多语言组合(mDenseOn 和 mLateOn)也以 +0.13 的差距复现了这一结果。规模也无法拯救单向量方案。Qwen3-Embedding-4B 是最强的稠密模型,其活跃(非嵌入)参数量约为我的模型的 33 倍,但仍落后 0.13,而 8B 版本的得分甚至低于 4B 版本。

  • WandbCallback:在已安装 wandb 的情况下,将训练指标记录到 W&B
  • TensorBoardCallback:在可访问 tensorboard 的情况下,将训练指标记录到 TensorBoard
  • CodeCarbonCallback:在已安装 codecarbon 的情况下,跟踪训练期间的碳排放

有关最实用的训练参数的更多信息,请参阅 Multi-Vector Encoder > Training Overview > Training Arguments。为了充分利用你的训练,值得一读。

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

训练器

全领域一览,按得分排序,并按架构家族着色。

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)

这就是经典的 ColBERT 流程:一个生成上下文相关 token 嵌入向量的 Transformer,一个将每个 token 嵌入向量降维至 128 维的 token 级 Dense 层,一个决定评分时哪些 token 计入的 MultiVectorMask 层,以及一个 token 级 Normalize 层。投影层从随机状态开始,因此该模型需要经过训练才能发挥作用。有趣的是,这种方法对强大的稠密嵌入骨干模型同样有效。在我的实验中,仅凭一个全新的投影层和 25k 个训练样本对,在 Alibaba-NLP/gte-modernbert-base 上从头训练的投影层就达到了与现有检查点起点相差 0.03 以内的性能。

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

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

最后,在“评估”部分,我将展示我微调得到的 multi-vector-encoder/mLateOn-medical 模型——它与本文同步在单张 RTX 3090 上训练了 14.5 小时——在我的医学检索评估中轻松超越了所有我能找到的通用检索模型,包括稠密、稀疏、词法以及多向量类型的模型。

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

这个结果让我很意外,而且在两个模型家族上都得到了复现。*无监督检查点在适应新领域方面的表现远好于它们已经完成训练的同类模型,尽管起点更低,却实现了反超。这些检查点位于大规模对比预训练之后、通用检索监督微调之前,因此它们保留了完整的后期交互结构,却没有那些领域训练随后不得不“抹掉”的通用调优痕迹。相比之下,已完成的检查点在我尝试的每一个学习率下都几乎没什么变化,甚至出现了倒退。

经典的 ColBERT token 化技巧([MASK] 查询扩展、[Q] / [D] 前缀 token、文档长度上限、标点跳过列表)默认全部关闭,且均可配置。完整列表请参阅“创建自定义模型”一节。顺带一提,我在领域微调时测试了四种配置的 [MASK] 查询扩展,均未产生可测量的差异,所以不必拘泥于经典配方。

© 版权声明
明晰AI

相关文章

广告也精彩

暂无评论

none
暂无评论...