ColBERT多向量预训练：突破信息检索新高度

妩媚怡口莲

1. ColBERT模型训练现状与挑战

ColBERT（Contextualized Late Interaction）这类多向量检索模型近年来在信息检索领域崭露头角，凭借其出色的跨领域泛化能力、长文本处理优势以及复杂推理检索特性，逐渐成为业界关注的焦点。与传统的单向量稠密检索模型相比，ColBERT通过为每个token生成独立的嵌入向量，再通过"延迟交互"机制计算query-document相关性，能够捕捉更细粒度的语义匹配关系。

然而当前主流实践存在一个明显的局限性：几乎所有state-of-the-art的ColBERT模型（包括我们之前发布的GTE-ModernColBERT和ColBERT-small）都是通过在预训练好的单向量模型基础上，简单添加一个知识蒸馏（KD）阶段来构建的。就连最新发布的mxbai-edge-colbert-v0模型，虽然详细描述了从零训练ColBERT的方法，但其早期阶段仍然完全在单向量设置下完成，仅在最后的蒸馏阶段才启用多向量目标函数。

这种做法的核心假设是：多向量特性只需要在模型训练的最后阶段引入即可。但这就引出了一个关键问题：如果我们从一开始就在多向量设置下进行预训练，会不会获得更好的效果？或者说，多向量预训练的成本是否值得，还是说最后的蒸馏步骤已经足够？

2. 模型训练的三阶段解析

要深入理解这个问题，我们需要拆解现代检索模型的典型训练流程。一个强大的检索模型通常需要经历三个关键训练阶段，每个阶段提供不同类型的学习信号：

2.1 无监督对比预训练

这是整个训练过程中计算成本最高的阶段。模型需要处理海量的文本对（查询语句和相关文档），仅使用批次内负样本（in-batch negatives）来学习区分相关与不相关文档。虽然不需要人工标注数据，但这个阶段通常需要数千GPU小时的计算资源。

在技术实现上，我们利用PyLate框架的两个关键特性来优化这个过程：

GradCache技术：突破GPU显存限制，实现任意大小的批次规模（标准梯度累积方法在对比学习中不适用）
多GPU样本聚合：在不增加内存需求的情况下，有效扩大批次规模

通过这些优化，我们可以轻松达到约16k的批次大小，这对于获得有意义的硬负样本至关重要。同时，PyLate集成的NanoBEIR工具让我们能够方便地进行训练中评估，极大地简化了超参数（如学习率和温度系数）的调优过程。

实践提示：使用split_batches选项确保每个批次仅包含单一数据源，这可以防止模型走捷径（shortcut learning），即仅通过数据来源而非实际内容来区分样本。

2.2 有监督对比微调

这是一个精炼阶段，使用规模较小但质量更高的数据集。这些数据集中包含了经过专门挖掘的硬负样本——那些表面相似但实际不相关的文档。与随机批次内负样本相比，这提供了更强的训练信号，因为模型需要学会区分真正相关文档与那些具有迷惑性的相似文档。

2.3 知识蒸馏（KD）

这个阶段将强大但计算昂贵的"教师模型"的知识迁移到我们的检索模型（"学生模型"）中。教师模型为每个查询语句生成一组文档的相关性分数，学生模型则学习复现这些分数分布。

3. ColBERT模型的规模化预训练

为了准确评估多向量预训练的效果，我们设计了严格的对照实验。我们选择使用Nomic Embed公开数据集进行训练，这有一个战略考量：Nomic已经基于相同的数据混合训练了一个稠密的ModernBERT模型，这让我们可以在数据、基础模型和训练流程完全一致的情况下，仅改变对比训练阶段是在稠密还是多向量设置下进行。

我们比较了三种训练流程来衡量每个阶段对最终性能的影响：

标准流程：在预训练好的稠密模型上仅进行KD步骤（设置a）
ColBERT-zero：所有预训练阶段都在多向量设置下完成（设置c）
混合流程：从Nomic的modernbert-embed-unsupervised稠密模型开始，仅在有监督和蒸馏步骤中使用多向量设置（设置b）

考虑到无监督阶段成本极高（约占有监督和KD步骤总和的10倍），我们还特别研究了是否可以通过某种方式"走捷径"——即跳过最昂贵的无监督预训练，同时仍能获得相似的性能。

4. BEIR基准测试结果分析

我们在BEIR基准上的评估揭示了一个重要发现：当前将多向量训练视为次要微调步骤的标准实践，实际上严重限制了模型的潜力。具体表现在：

4.1 完全预训练的优势

我们的完全多向量预训练模型ColBERT-Zero（仅使用公开Nomic数据训练）达到了55.43的nDCG@10分数，超越了之前state-of-the-art的GTE-ModernColBERT（54.67）。这个结果尤其值得注意，因为：

基于Nomic数据训练的稠密基线模型得分为52.89
基于GTE专有数据训练的稠密模型得分为55.33

这意味着，尽管初始数据质量存在2.4分的差距，完整的ColBERT预训练流程让ColBERT-Zero不仅弥补了这个差距，甚至超过了gte-modernbertbase模型，为<150M参数的模型在BEIR上设立了新的state-of-the-art。

4.2 仅使用KD的局限性

仅使用知识蒸馏的方法（流程a）得分为54.09，比完整预训练低了1.3分。这证实了简单的蒸馏步骤对于优化多向量性能是不够的。

4.3 有监督+KD的高效替代方案

我们发现，通过在有监督对比步骤就切换到多向量设置，再进行蒸馏（流程b），可以获得55.12的nDCG@10分数，几乎达到了完整预训练模型（55.43）的性能，而计算成本仅为约40 GH200小时，相比完整预训练的408小时，实现了10倍的成本节约，同时保留了99.4%的性能。

5. 提示对齐的关键作用

Nomic的基础模型在预训练时使用了特定的非对称提示词（"search_query:"和"search_document:"）作为查询和文档的前缀。虽然ColBERT本身已有通过[Q]和[D]标记实现的非对称机制，但我们发现：

在微调时去除这些预训练提示词会导致显著的性能下降
相反，给没有使用提示词预训练的模型添加提示词同样会损害性能

关键结论是：无论你的基础模型使用何种提示词设置，在下游训练中都必须保持一致。但提示对齐并非全部原因——即使在完全对齐的情况下（比较使用和不使用提示词的完整ColBERT预训练），我们仍观察到明显的性能差距（54.61 vs. 55.43 nDCG@10），这表明提示词本身具有内在优势。

我们的主要假设是：提示词标记充当了隐式的查询扩展——这些不携带特定含义的额外位置让模型能够存储关于序列的全局信息。原始ColBERT模型使用PAD标记实现类似目的，但现代Flash Attention实现禁用了这个技巧，因为它们不会为掩码标记生成可用嵌入。显式提示词可能在不经意间重新启用了这一机制。