Training and Finetuning Multi-Vector Embedding Models with Sentence Transformers

查看原文 / View Original 编辑 / Edit

Training and Finetuning Multi-Vector Embedding Models with Sentence Transformers

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

Sentence Transformers is a Python library for using and training embedding and reranker models for a wide range of applications, such as retrieval augmented generation, semantic search, semantic textual similarity, and more. Its v6.0 update introduces a fourth model type: MultiVectorEncoder, for ColBERT-style late interaction retrieval, alongside a complete training approach for it. In this blogpost, I’ll show you how to use it to finetune a multi-vector model that outperforms general-purpose retrievers on your data. This method can also train strong new multi-vector models from scratch. Everything below runs on pip install -U "sentence-transformers[train]".

Sentence Transformers 是一个用于使用和训练嵌入（embedding）及重排序（reranker）模型的 Python 库，广泛应用于检索增强生成（RAG）、语义搜索、语义文本相似度等领域。其 v6.0 版本引入了第四种模型类型：MultiVectorEncoder，用于 ColBERT 风格的后期交互（late interaction）检索，并提供了完整的训练方案。在这篇博文中，我将展示如何使用它来微调多向量模型，使其在你的数据上表现优于通用检索器。该方法也可以从零开始训练强大的多向量模型。以下所有内容均可通过 pip install -U "sentence-transformers[train]" 运行。

Finetuning multi-vector models involves several components: the model itself, datasets, loss functions, training arguments, evaluators, and the trainer class. I’ll have a look at each of these components, accompanied by practical examples of how they can be used for finetuning strong multi-vector models. Lastly, in the Evaluation section, I’ll show you that my finetuned multi-vector-encoder/mLateOn-medical model, trained in 14.5 hours on a single RTX 3090 alongside this blogpost, easily outperforms every general-purpose retrieval model I could find on my medical retrieval evaluation: dense, sparse, lexical, and multi-vector alike.

微调多向量模型涉及多个组件：模型本身、数据集、损失函数、训练参数、评估器以及训练器类。我将逐一介绍这些组件，并附带如何使用它们微调强大模型的可操作示例。最后，在“评估”部分，我将展示我微调的 multi-vector-encoder/mLateOn-medical 模型——它是在撰写本文期间使用单张 RTX 3090 显卡耗时 14.5 小时训练完成的——在我的医学检索评估中，轻松超越了我所能找到的所有通用检索模型，无论是稠密（dense）、稀疏（sparse）、词法（lexical）还是多向量模型。

If you’re interested in finetuning dense embedding models, sparse embedding models, or rerankers instead, then consider reading through my prior Training and Finetuning Embedding Models, Training and Finetuning Sparse Embedding Models, and Training and Finetuning Reranker Models blogposts. This blogpost is about training multi-vector models. If you want to learn how to use them, from loading and encoding to indexing in vector databases, see the companion Multi-Vector (Late Interaction) Embedding Models with Sentence Transformers blogpost.

如果你对微调稠密嵌入模型、稀疏嵌入模型或重排序模型感兴趣，请阅读我之前的博文：《训练和微调嵌入模型》、《训练和微调稀疏嵌入模型》以及《训练和微调重排序模型》。本篇博文专注于多向量模型的训练。如果你想了解如何使用它们（从加载、编码到在向量数据库中建立索引），请参阅配套博文《使用 Sentence Transformers 的多向量（后期交互）嵌入模型》。

What are Multi-Vector models?

什么是多向量模型？

A dense embedding model compresses a whole text into a single vector, and similarity is one dot product between two such summaries. A multi-vector model (also called a late-interaction or ColBERT-style model) skips that compression. It keeps one small vector per token and scores a query against a document with the MaxSim operator, where every query token finds its best-matching document token and the scores are summed. Token-level matching preserves exactly the fine-grained signals that a single vector has to average away, which usually means stronger retrieval, at the cost of a bigger index.

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

Why Finetune?

为什么要微调？

Finetuning multi-vector models significantly improves their retrieval performance on your specific domain: the vocabulary, the query style, and the notion of relevance all differ between web search, legal discovery, code search, and scientific literature review. Because queries and documents are matched token by token, multi-vector models pick up fine-grained domain signals that single-vector models tend to average away, and they respond very well to even modest amounts of in-domain finetuning data.

微调多向量模型可以显著提升其在特定领域内的检索性能：词汇表、查询风格以及相关性的定义在网络搜索、法律发现、代码搜索和科学文献综述中各不相同。由于查询和文档是按 token 进行匹配的，多向量模型能够捕捉到单向量模型倾向于平均掉的细粒度领域信号，并且即使只有少量的领域内微调数据，它们也能表现出极佳的响应。

Beyond that, most released retrieval models were configured for short passages. The classic ColBERT checkpoints truncate documents at 180 or 300 tokens, and many popular dense models at 256 or 512, because their MS MARCO-style training data rarely goes beyond that. If your documents are long, these models silently discard most of every document before scoring it. On my medical evaluation with passages averaging 941 tokens, I measured that this truncation costs up to 0.24 NDCG@10, considerably more than any difference between model architectures.

此外，大多数已发布的检索模型都是针对短文本配置的。经典的 ColBERT 检查点将文档截断为 180 或 300 个 token，许多流行的稠密模型则截断为 256 或 512 个，因为它们的 MS MARCO 风格训练数据很少超过这个长度。如果你的文档很长，这些模型在评分前会默默丢弃文档的大部分内容。在我的医学评估中（段落平均长度为 941 个 token），我测得这种截断会导致高达 0.24 的 NDCG@10 损失，这远超模型架构差异带来的影响。

Training Components

训练组件

Training MultiVectorEncoder models involves the following components:

  • Model: The model to finetune or the architecture to build fresh.
  • Dataset: The data used for training and evaluation.
  • Loss Function: A function that measures the model’s performance and guides the optimization process.
  • Training Arguments (optional): Parameters that impact training performance, tracking, and debugging.
  • Evaluator (optional): A class for evaluating the model before, during, or after training.
  • Trainer: Brings together all training components.

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

  • 模型 (Model)：要微调的模型或从零构建的架构。
  • 数据集 (Dataset)：用于训练和评估的数据。
  • 损失函数 (Loss Function)：衡量模型性能并指导优化过程的函数。
  • 训练参数 (Training Arguments, 可选)：影响训练性能、跟踪和调试的参数。
  • 评估器 (Evaluator, 可选)：用于在训练前、训练中或训练后评估模型的类。
  • 训练器 (Trainer)：整合所有训练组件。

Model

模型

Multi-vector training gives you a real choice of starting point, and it matters more than you might expect.

Finetuning an existing multi-vector model

多向量训练为你提供了真正的起点选择，这比你预想的要重要得多。

微调现有的多向量模型

If you want to further finetune an existing multi-vector model, you don’t have to worry about the architecture at all: 如果你想进一步微调现有的多向量模型，完全不必担心架构问题：

from sentence_transformers import MultiVectorEncoder

# Loading in fp32 is preferred for training if your memory can handle it
# 如果内存允许，训练时建议使用 fp32 加载
model = MultiVectorEncoder(
    "lightonai/mLateOn-unsupervised",
    model_kwargs={"torch_dtype": "float32"},
    processor_kwargs={"model_max_length": 8192}, # the tokenizer-level token limit
)
← 返回新闻列表 / Back to news list