多向量模型微调实战：Sentence Transformers 构建医疗检索冠军

1 阅读

在信息检索领域，传统的稠密嵌入模型（Dense Embedding Models）虽然训练高效且推理迅速，但在处理长文档和细粒度语义匹配时往往力不从心。它们将整段文本压缩为单一向量，不可避免地丢失了关键的分词信号。相比之下，多向量模型（Multi-Vector Models），也被称为晚交互模型（Late-Interaction Models）或 ColBERT 风格模型，通过为每个 token 生成独立向量并利用 MaxSim 算子进行打分，保留了丰富的细粒度信息。然而，这些模型通常计算开销较大，且通用预训练模型在特定垂直领域（如医疗、法律、代码）的表现往往不尽如人意。

本文旨在深入探讨如何利用 Sentence Transformers 库，从零开始或基于现有检查点，高效微调多向量嵌入模型。我们将通过一个具体的医疗检索案例，展示如何在单张消费级 GPU 上，仅用 14.5 小时训练出一个超越主流通用检索模型的性能。这不仅涉及技术细节的堆砌，更包含了对训练策略、损失函数选择、评估指标优化以及索引压缩等关键问题的深度剖析。

多向量模型的核心优势与微调必要性

多向量模型的核心机制在于其独特的评分方式。不同于稠密模型计算两个文档向量的点积，多向量模型将查询和文档分别 token 化，为每个 token 生成一个低维向量（通常为 128 维），然后在查询和文档之间进行“晚交互”（Late Interaction）。具体而言，对于查询中的每一个 token，模型会在文档中找到最相似的 token 向量，计算其相似度，并将所有匹配对的相似度求和作为最终得分。这种机制完美保留了 token 级别的语义对齐能力，使得模型能够捕捉到那些被平均化操作抹去的细微差别。

然而，这种优势也伴随着挑战。首先，多向量模型的索引规模远大于稠密模型。存储每个 token 的向量意味着索引体积可能扩大数十倍。其次，通用预训练模型（如基于 MS MARCO 训练的检查点）通常针对短文本优化，文档长度截断在 180 到 512 个 token 之间。当面对平均长度超过 900 个 token 的长文档时，这些模型会静默丢弃大部分信息，导致检索性能大幅下降。

微调（Finetuning）在此刻显得尤为关键。通过引入领域特定的数据，模型可以学习到该领域的词汇偏好、查询风格以及相关性定义。例如，在医疗领域，术语的精确匹配和上下文逻辑至关重要，通用模型难以掌握这些细微的领域信号。此外，微调允许我们自定义文档长度限制，确保模型能够充分利用长文档中的全部信息。LightOn 团队在代码检索领域的实践也证明了这一点，他们通过微调 LateOn 模型构建了 LateOn-Code，显著提升了特定场景下的检索效果。

训练组件详解：从模型架构到评估策略

构建一个强大的多向量模型训练系统，需要精心配置多个核心组件。这些组件共同协作，确保训练过程既高效又有效。

模型选择与初始化

多向量训练提供了灵活的起点选择。如果你希望进一步微调现有的多向量模型，可以直接加载预训练检查点，如 lightonai/mLateOn-unsupervised。该检查点已经过大规模对比预训练，具备完整的晚交互结构，但尚未针对特定领域进行监督微调。

加载模型时，需要注意几个关键参数。首先，建议使用 float32 精度进行训练，以平衡精度与显存占用。其次，必须调整模型内部的长度限制。许多通用检查点限制了查询和文档的最大长度，这会导致长文档被截断。在微调前，应将 query_length 和 document_length 设置为 None，让模型根据 tokenizer 的 model_max_length 进行动态截断，从而适应你的数据分布。

此外，还可以引入标点符号跳过列表（Skiplist）。在文档侧的评分中排除标点符号，不仅能减少索引大小（在我的案例中减少了 9.6%），还能提升检索质量，因为标点符号通常不携带语义信息。

如果你没有现成的多向量检查点，也可以从基础 Transformer 模型开始构建。MultiVectorEncoder 支持将任意基础模型（如 Alibaba-NLP/gte-modernbert-base）作为起点，并自动附加一个随机的 token 级投影头。虽然从头开始训练需要更多数据，但在资源有限或需要完全定制化架构时，这是一个可行的选择。

数据集构建与格式

数据集是微调的灵魂。对于多向量模型，最简单且有效的数据格式是（查询，相关文档）对。在 Hugging Face Hub 上，有许多现成的数据集可供使用，如 MIRIAD，它包含了数百万个医疗问题及其来源段落。你也可以使用 datasets 库加载本地数据，支持 CSV、JSON、Parquet 等多种格式。

数据集的列顺序至关重要。第一个列被视为查询，后续列被视为文档。如果涉及知识蒸馏，数据格式可能需要调整为包含多个候选文档及其教师模型评分的形式。确保数据集的规模适中，既要有足够的样本量以覆盖领域特征，又要避免过拟合。在我的实验中，100 万对样本在单卡上训练了 14.5 小时，而 10 万对样本仅需 75 分钟，且性能差距较小，这表明在资源受限的情况下，小规模高质量数据同样有效。

损失函数配置

选择合适的损失函数是训练成功的关键。对于最常见的问答对任务，MultiVectorMultipleNegativesRankingLoss 是标准选择。该损失函数在批次内为每个查询生成负样本，通过对比学习优化排序能力。

然而，多向量模型有一个特殊的陷阱：对比损失的默认缩放因子（scale）为 1.0，而稠密模型通常为 20.0。这是因为多向量模型的 MaxSim 分数范围是 [0, 查询长度]，而稠密模型的余弦相似度范围是 [-1, 1]。如果错误地沿用稠密模型的 scale 值，会导致梯度饱和，训练失效。因此，必须使用 scale=1.0 或根据具体情况进行调整。

为了缓解显存压力，可以使用 CachedMultiVectorMultipleNegativesRankingLoss。该变体通过梯度缓存（GradCache）技术，将有效批次大小与显存占用解耦。你可以设置较小的 mini_batch_size（如 16）来限制显存，同时保持较大的有效批次大小（如 128），从而在不增加显存消耗的前提下提升训练效率。

训练参数与优化器

训练参数的设置直接影响收敛速度和最终性能。学习率（Learning Rate）是一个需要精细调优的参数。在我的实验中，1e-4 的学习率表现最佳，高于稠密模型常用的 2e-5 或 5e-5。这可能是因为多向量模型需要更强烈的梯度信号来对齐 token 级别的向量。

此外，必须正确配置提示词（Prompts）。训练时，模型不会自动应用检查点中存储的 [Q] 和 [D] 标记，你需要显式地在训练参数中指定这些标记对应的列名，以确保训练时的输入格式与推理时一致。

评估与回调

评估是指导训练过程的重要手段。MultiVectorInformationRetrievalEvaluator 允许你使用自定义的查询集和文档集来评估模型性能。为了获得有意义的评估结果，评估集应该具有足够的难度，避免模型轻易得分过高。如果评估集过于简单，可以引入干扰文档（Distractors）来增加区分度。

回调函数（Callbacks）提供了额外的监控能力。你可以集成 WandbCallback 或 TensorBoardCallback 来记录训练指标，使用 CodeCarbonCallback 追踪碳足迹。这些工具不仅有助于调试，还能帮助你在资源受限的环境中做出更明智的决策。

实战案例：医疗检索模型的微调与评估

为了验证上述策略的有效性，我构建了一个名为 mLateOn-medical 的医疗检索模型。该模型基于 lightonai/mLateOn-unsupervised 检查点，在 MIRIAD 数据集上进行了微调。

训练流程

训练过程遵循了以下关键步骤：

  1. 加载模型：使用 float32 精度加载 mLateOn-unsupervised 检查点，并移除长度限制。
  2. 预处理：添加标点符号跳过列表，优化索引效率。
  3. 数据加载：从 MIRIAD 数据集加载 100 万对医疗问题与段落。
  4. 损失函数：配置 CachedMultiVectorMultipleNegativesRankingLoss，设置 mini_batch_size=16。
  5. 训练参数：设置学习率为 1e-4，批次大小为 128，启用 bf16 加速。
  6. 评估：使用 500 个 held-out 问题作为开发集，监控 NDCG@10 指标。

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

整个训练过程在单张 RTX 3090 上运行，峰值显存占用为 17.5 GB。训练完成后，模型在 1000 个测试问题上的 NDCG@10 达到了 0.9398，相比零-shot 基线提升了 0.0311。

性能对比与结果分析

为了全面评估模型效果，我将 mLateOn-medical 与超过 50 种其他检索模型进行了对比，涵盖了稠密、稀疏、多向量以及基于 BM25 的算法。评估集包含 1000 个医疗问题，在 20 万个唯一段落中进行检索。

结果显示，mLateOn-medical 以 0.9139 的 NDCG@10 位居榜首，显著优于最强的零-shot 多向量模型 lightonai/mLateOn（0.8520）和稠密模型 Qwen/Qwen3-Embedding-4B（0.7817）。这一差距表明，领域特定的微调能够带来巨大的性能提升。

值得注意的是，即使将通用多向量模型的长度限制提升至 1024 token，其性能也无法与微调后的模型相提并论。这进一步证明了微调在适应长文档和领域特征方面的不可替代性。

索引优化策略

多向量模型的主要劣势在于索引体积。对于 20 万个段落，原始多向量索引需要约 45 GB 的存储空间。为了缓解这一问题，我采用了多种优化策略：

  1. 分层 Token 池化（Hierarchical Token Pooling）：通过聚类 token 向量并存储聚类中心，可以将索引大小减少至原来的四分之一，同时保持 0.8991 的 NDCG@10。
  2. 残差量化（Residual Quantization）：结合 fast-plaid 库，使用 1-bit 残差量化和文档侧剪枝，可以将索引压缩至 1.45 GB，同时保持 0.8642 的 NDCG@10。

这一优化策略将多向量索引的大小缩小了 30 倍，使其在存储成本上接近甚至低于某些稠密模型，证明了多向量检索在大规模部署中的可行性。

总结与展望

本文详细阐述了使用 Sentence Transformers 微调多向量嵌入模型的全过程。通过精心选择模型起点、优化损失函数配置、合理设置训练参数以及实施索引压缩策略，我们成功构建了一个在医疗领域表现卓越的检索模型。

多向量模型虽然面临索引规模大的挑战，但通过合理的工程优化，这一劣势可以被有效克服。其保留的细粒度语义信息使其在处理长文档和复杂查询时具有独特优势。对于任何希望在特定领域构建高性能检索系统的开发者来说，掌握多向量模型的微调技术都是必不可少的技能。

未来，随着多向量模型在更多领域的应用，我们期待看到更多针对垂直领域的专用模型涌现。同时，索引压缩技术的进步也将进一步降低多向量检索的部署门槛，使其成为通用检索系统的重要补充。通过持续的技术创新与实践探索，多向量模型必将在信息检索领域发挥更大的作用。