
更新:在 2024 年 8 月 31 日,我们发布了 Jina-ColBERT 的第二个版本,具有更好的性能、支持 89 种语言以及灵活的输出维度。查看发布文章了解更多详情。
上周五,Jina AI 在 Hugging Face 上发布的 ColBERT 模型在 AI 社区引起了巨大轰动,尤其是在 Twitter/X 上。虽然许多人都熟悉开创性的 BERT 模型,但 ColBERT 引发的热议让一些人产生疑问:在信息检索技术领域,是什么让 ColBERT 如此出众?为什么 AI 社区对 8192 长度的 ColBERT 如此兴奋?本文将深入探讨 ColBERT 和 ColBERTv2 的细节,重点介绍它们的设计、改进以及 ColBERT 后期交互的惊人效果。

tag什么是 ColBERT?
"ColBERT" 代表"基于 BERT 的上下文化后期交互"(Contextualized Late Interaction over BERT),这是一个源自斯坦福大学的模型,它利用了 BERT 的深度语言理解能力,同时引入了一种新颖的交互机制。这种被称为"后期交互"的机制,通过在检索过程的最后阶段之前分别处理查询和文档,实现了高效和精确的检索。具体来说,该模型有两个版本:
- ColBERT:最初的模型是 Omar Khattab 和 Matei Zaharia 的创意,通过论文"ColBERT:通过基于 BERT 的上下文化后期交互实现高效和有效的段落搜索"提出了一种新颖的信息检索方法。该论文发表于 SIGIR 2020。

原始 ColBERT 论文介绍了"后期交互"。
- ColBERTv2:在基础工作的基础上,Omar Khattab 继续他的研究,与 Barlas Oguz、Matei Zaharia 和 Michael S. Bernstein 合作,在 SIGIR 2021 上提出了"ColBERTv2:通过轻量级后期交互实现有效和高效的检索"。ColBERT 的这一新版本解决了之前的局限性,并引入了关键改进,如去噪监督和残差压缩,提高了模型的检索效果和存储效率。

ColBERTv2 添加了去噪监督和残差压缩,以提高训练数据质量并减少空间占用。
tag理解 ColBERT 的设计
鉴于 ColBERTv2 的架构与原始 ColBERT 非常相似,其主要创新在于训练技术和压缩机制,我们首先将深入探讨原始 ColBERT 的基本方面。
tag什么是 ColBERT 中的后期交互?
"交互"指的是通过比较查询和文档的表示来评估它们之间相关性的过程。
"后期交互"是 ColBERT 的精髓。这个术语源自模型的架构和处理策略,其中查询和文档表示之间的交互发生在过程的后期,即在两者被独立编码之后。这与"早期交互"模型形成对比,在早期交互模型中,查询和文档嵌入在较早的阶段就进行交互,通常是在模型编码之前或期间。
交互类型 | 模型 |
---|---|
早期交互 | BERT, ANCE, DPR, Sentence-BERT, DRMM, KNRM, Conv-KNRM, etc. |
后期交互 | ColBERT, ColBERTv2 |
早期交互会增加计算复杂度,因为它需要考虑所有可能的查询-文档对,这使得它在大规模应用中效率较低。
像 ColBERT 这样的后期交互模型通过允许预计算文档表示并在最后采用更轻量级的交互步骤来优化效率和可扩展性,这种交互步骤专注于已编码的表示。这种设计选择使得检索时间更快,计算需求更低,更适合处理大型文档集合。
tag无交互:文档和查询嵌入向量的余弦相似度
许多实用的向量数据库和神经搜索解决方案都依赖于文档和查询嵌入向量之间的快速余弦相似度匹配。虽然这种方法在直观性和计算效率方面很有吸引力,但这种被称为"无交互"或"非基于交互"的方法与那些包含某种形式的查询和文档交互的模型相比,已被发现表现不佳。
"无交互"方法的核心限制在于其无法捕捉查询和文档术语之间的复杂细微差别和关系。信息检索的本质是理解并匹配查询背后的意图与文档中的内容。这个过程通常需要对相关术语进行深入的上下文理解,而单一的、聚合的文档和查询嵌入向量难以提供这种理解。
tagColBERT 中的查询和文档编码器
ColBERT 的编码策略基于以深度上下文语言理解而闻名的 BERT 模型。该模型为查询或文档中的每个 token 生成密集的向量表示,分别为查询和文档创建上下文化嵌入向量包。这为后期交互阶段的嵌入向量细致比较提供了基础。
tagColBERT 的查询编码器
对于包含 token 的查询 ,首先将 分词为基于 BERT 的 WordPiece token,并在前面添加特殊的 [Q]
token。这个 [Q]
token 位于 BERT 的 [CLS]
token 之后,表示查询的开始。
如果查询短于预定义的 token 数量 ,则用 [mask]
token 填充到 ;否则截断为前 个 token。填充后的序列通过 BERT,然后经过 CNN(卷积神经网络)和标准化。输出是一组嵌入向量,在下面表示为 :
tagColBERT 的文档编码器
类似地,对于包含 token 的文档 ,在前面添加 [D]
token 来表示文档的开始。这个序列无需填充,经过相同的处理过程,得到的嵌入向量集合表示为 :
使用 [mask]
token 填充查询(在论文中称为查询增强)确保了所有查询的长度一致,便于批处理。[Q]
和 [D]
token 明确标记查询和文档的开始,帮助模型区分这两种输入类型。
tagColBERT 与交叉编码器的比较
交叉编码器同时处理查询和文档对,使其具有很高的准确性,但由于需要评估每个可能的配对,在大规模任务中效率较低。它们在特定场景中表现出色,比如语义相似度任务或详细的内容比较,这些场景需要精确评分句子对。然而,这种设计限制了它们在需要从大型数据集中快速检索的情况下的应用,这种情况下预计算嵌入向量和高效的相似度计算至关重要。
相比之下,ColBERT 的后期交互模型允许预计算文档嵌入向量,在不影响语义分析深度的同时显著加快了检索过程。虽然与交叉编码器的直接方法相比这种方法看似反直觉,但它为实时和大规模信息检索任务提供了一个可扩展的解决方案。这代表了计算效率和交互建模质量之间的战略性妥协。
tag使用 ColBERT 查找 top-K 文档
一旦我们获得了查询和文档的嵌入向量,查找最相关的 top-K 文档就变得直接了(但不像计算两个向量的余弦相似度那么直接)。
关键操作包括批量点积计算以计算词项间的相似度,在文档词项上进行最大池化以找到每个查询词项的最高相似度,以及对查询词项求和以得到总文档分数,然后根据这些分数对文档进行排序。伪 PyTorch 代码如下所示:
import torch
def compute_relevance_scores(query_embeddings, document_embeddings, k):
"""
Compute relevance scores for top-k documents given a query.
:param query_embeddings: Tensor representing the query embeddings, shape: [num_query_terms, embedding_dim]
:param document_embeddings: Tensor representing embeddings for k documents, shape: [k, max_doc_length, embedding_dim]
:param k: Number of top documents to re-rank
:return: Sorted document indices based on their relevance scores
"""
# Ensure document_embeddings is a 3D tensor: [k, max_doc_length, embedding_dim]
# Pad the k documents to their maximum length for batch operations
# Note: Assuming document_embeddings is already padded and moved to GPU
# Compute batch dot-product of Eq (query embeddings) and D (document embeddings)
# Resulting shape: [k, num_query_terms, max_doc_length]
scores = torch.matmul(query_embeddings.unsqueeze(0), document_embeddings.transpose(1, 2))
# Apply max-pooling across document terms (dim=2) to find the max similarity per query term
# Shape after max-pool: [k, num_query_terms]
max_scores_per_query_term = scores.max(dim=2).values
# Sum the scores across query terms to get the total score for each document
# Shape after sum: [k]
total_scores = max_scores_per_query_term.sum(dim=1)
# Sort the documents based on their total scores
sorted_indices = total_scores.argsort(descending=True)
return sorted_indices
注意,这个过程在训练和推理时的重排序中都会使用。ColBERT 模型使用成对排序损失进行训练,其中训练数据由三元组 组成,其中 表示查询, 是查询的相关(正面)文档, 是不相关(负面)文档。模型的目标是学习表示,使得 和 之间的相似度分数高于 和 之间的分数。
训练目标可以数学表示为最小化以下损失函数:
,其中 表示 ColBERT 计算的查询 和文档 之间的相似度分数。这个分数是通过聚合查询和文档之间最佳匹配嵌入向量的最大相似度分数获得的,遵循模型架构中描述的后期交互模式。这种方法确保模型被训练为区分给定查询的相关和不相关文档,通过鼓励正面和负面文档对的相似度分数之间有更大的差距。
tagColBERTv2 中的去噪监督
ColBERTv2 中的去噪监督通过选择具有挑战性的负样本和利用交叉编码器进行蒸馏来改进原始训练过程。这种改进训练数据质量的复杂方法包括以下几个步骤:
- 初始训练:利用 MS MARCO 数据集的官方三元组,包括查询、相关文档和不相关文档。
- 索引和检索:使用 ColBERTv2 的压缩方法索引训练段落,然后为每个查询检索 top-k 段落。
- 交叉编码器重排序:通过 MiniLM 交叉编码器重排序来改进段落选择,将其分数蒸馏到 ColBERTv2 中。
- 形成训练元组:生成 w-way 元组用于训练,结合高排名和低排名的段落创建具有挑战性的样本。
- 迭代改进:重复该过程以持续改进硬负样本的选择,从而提高模型性能。
注意,这个过程代表了对 ColBERT 训练机制的复杂改进,而不是对其架构的根本改变。
tagColBERT 的超参数
ColBERT 的超参数总结如下:
超参数 | 最佳选择 | 原因 |
---|---|---|
学习率 | 3 x 10^{-6} | 为微调选择,确保模型更新稳定有效。 |
批量大小 | 32 | 平衡计算效率和每次更新捕获足够信息的能力。 |
每个查询的嵌入数量 (Nq) | 32 | 固定值以确保查询间表示大小的一致性,有助于高效处理。 |
嵌入维度 (m) | 128 | 经证明可以在表示能力和计算效率之间提供良好平衡。 |
训练迭代次数 | 200k (MS MARCO),125k (TREC CAR) | 选择以确保充分学习同时避免过拟合,根据数据集特征进行调整。 |
每维度嵌入的字节数 | 4(重排序),2(端到端排序) | 在精度和空间效率之间权衡,考虑应用场景(重排序与端到端)。 |
向量相似度函数 | 余弦(重排序),(平方)L2(端到端) | 基于各自检索场景的性能和效率选择。 |
FAISS 索引分区 (P) | 2000 | 决定搜索空间分区的粒度,影响搜索效率。 |
搜索最近分区数 (p) | 10 | 在搜索范围和计算效率之间取得平衡。 |
每个嵌入的子向量数 (s) | 16 | 影响量化的粒度,影响搜索速度和内存使用。 |
索引每维度表示 | 16 位值 | 用于端到端检索的第二阶段,在准确性和空间之间取得平衡。 |
编码器层数 | 12 层 BERT | 上下文理解深度和计算效率之间的最佳平衡。 |
最大查询长度 | 128 | 查询编码器处理的最大 token 数。这在 Jina-ColBERT 模型中得到了扩展。 |
最大文档长度 | 512 | 文档编码器处理的最大 token 数。这在 Jina-ColBERT 模型中扩展到了 8192。 |
tagColBERT 的索引策略
与将每个文档编码为一个嵌入向量的基于表示的方法不同,ColBERT 将文档(和查询)编码为嵌入包,文档中的每个 token 都有自己的嵌入。这种方法本质上意味着对于较长的文档,需要存储更多的嵌入,这是原始 ColBERT 的一个痛点,后来在 ColBERTv2 中得到了解决。
高效管理这一问题的关键在于 ColBERT 使用向量数据库(如 FAISS)进行索引和检索,以及其精心设计的索引过程,可以高效处理大量数据。原始 ColBERT 论文提到了几个提高索引和检索效率的策略,包括:
- 离线索引:文档表示在离线计算,允许预先计算和存储文档嵌入。这个过程利用批处理和 GPU 加速来高效处理大型文档集合。
- 嵌入存储:文档嵌入可以使用 32 位或 16 位值存储每个维度,在精度和存储需求之间提供权衡。这种灵活性使 ColBERT 能够在有效性(检索性能)和效率(存储和计算成本)之间保持平衡。
在 ColBERTv2 中引入的残差压缩是原始 ColBERT 中不存在的新方法,它在保持质量的同时将模型的空间占用减少了 6-10 倍。这种技术通过有效捕获和仅存储与固定参考质心的差异来进一步压缩嵌入。
tagColBERT 的有效性和效率
人们可能最初认为将 BERT 的深度上下文理解整合到搜索中必然需要大量计算资源,由于高延迟和计算成本而使这种方法在实时应用中不太可行。然而,ColBERT 通过其创新的后期交互机制挑战并推翻了这一假设。以下是一些值得注意的要点:
- 显著的效率提升:与传统的基于 BERT 的排序模型相比,ColBERT 在计算成本(FLOPs)和延迟方面实现了数量级的减少。具体来说,对于给定的模型大小(例如 12 层"base" transformer 编码器),ColBERT 不仅匹配而且在某些情况下超过了 BERT 基础模型的效果,同时显著降低了计算需求。例如,在重排序深度 k=10 时,BERT 需要的 FLOPs 比 ColBERT 多近 180 倍;随着 k 的增加,这个差距进一步扩大,在 k=1000 时达到 13900 倍,在 k=2000 时甚至达到 23000 倍。
- 端到端检索中改进的召回率和 MRR@10:与最初认为查询和文档表示之间需要更深层的交互(如早期交互模型所见)才能获得高检索性能的直觉相反,ColBERT 的端到端检索设置展现出了卓越的效果。例如,其 Recall@50 超过了官方 BM25 的 Recall@1000 和几乎所有其他模型的 Recall@200,突显了该模型无需直接比较每个查询-文档对就能从庞大集合中检索相关文档的卓越能力。
- 实际应用的实用性:实验结果突出了 ColBERT 在实际场景中的实用性。其索引吞吐量和内存效率使其能够在几小时内索引像 MS MARCO 这样的大型文档集合,同时保持高效性和可管理的空间占用。这些特性突显了 ColBERT 适合部署在性能和计算效率都至关重要的生产环境中。
- 文档集合规模的可扩展性:也许最令人惊讶的结论是 ColBERT 在处理大规模文档集合时的可扩展性和效率。该架构允许预计算文档嵌入并利用高效的批处理进行查询-文档交互,使系统能够随文档集合规模有效扩展。考虑到有效文档检索所需的复杂性和理解深度,这种可扩展性是反直觉的,展示了 ColBERT 在平衡计算效率和检索效果方面的创新方法。
tag使用 jina-colbert-v1-en:一个支持 8192 长度的 ColBERTv2 模型
Jina-ColBERT 旨在实现快速且准确的检索,支持长达 8192 的上下文长度,利用了 JinaBERT 的进步,其架构增强使其能够处理更长的序列。
[D],[CLS]
填充。
tagJina 对原始 ColBERT 的改进
Jina-ColBERT 的主要进步在于其骨干网络 jina-bert-v2-base-en
,它能够处理显著更长的上下文(最多 8192 个 token),相比使用 bert-base-uncased
的原始 ColBERT 有了很大提升。这种能力对于处理包含大量内容的文档至关重要,可以提供更详细和更具上下文的搜索结果。
tagjina-colbert-v1-en 与 ColBERTv2 的性能比较
我们在 BEIR 数据集和偏好长上下文的新 LoCo 基准上评估了 jina-colbert-v1-en,将其与原始 ColBERTv2 实现和非交互式的jina-embeddings-v2-base-en 模型。
Dataset | ColBERTv2 | jina-colbert-v1-en | jina-embeddings-v2-base-en |
---|---|---|---|
Arguana | 46.5 | 49.4 | 44.0 |
Climate-Fever | 18.1 | 19.6 | 23.5 |
DBPedia | 45.2 | 41.3 | 35.1 |
FEVER | 78.8 | 79.5 | 72.3 |
FiQA | 35.4 | 36.8 | 41.6 |
HotpotQA | 67.5 | 65.9 | 61.4 |
NFCorpus | 33.7 | 33.8 | 32.5 |
NQ | 56.1 | 54.9 | 60.4 |
Quora | 85.5 | 82.3 | 88.2 |
SCIDOCS | 15.4 | 16.9 | 19.9 |
SciFact | 68.9 | 70.1 | 66.7 |
TREC-COVID | 72.6 | 75.0 | 65.9 |
Webis-touch2020 | 26.0 | 27.0 | 26.2 |
LoCo | 74.3 | 83.7 | 85.4 |
Average | 51.7 | 52.6 | 51.6 |
这个表格展示了 jina-colbert-v1-en 的出色性能,特别是在需要更长上下文长度的场景中,相比原始的 ColBERTv2。注意,jina-embeddings-v2-base-en 使用了更多的训练数据,而 jina-colbert-v1-en 仅使用 MSMARCO,这可能解释了为什么 jina-embeddings-v2-base-en 在某些任务上表现良好。
tagjina-colbert-v1-en 的使用示例
这段代码展示了使用 Jina-ColBERT 的索引过程,展示了其对长文档的支持。
from colbert import Indexer
from colbert.infra import Run, RunConfig, ColBERTConfig
n_gpu: int = 1 # Set your number of available GPUs
experiment: str = "" # Name of the folder where the logs and created indices will be stored
index_name: str = "" # The name of your index, i.e. the name of your vector database
if __name__ == "__main__":
with Run().context(RunConfig(nranks=n_gpu, experiment=experiment)):
config = ColBERTConfig(
doc_maxlen=8192 # Our model supports 8k context length for indexing long documents
)
indexer = Indexer(
checkpoint="jinaai/jina-colbert-v1-en",
config=config,
)
documents = [
"ColBERT is an efficient and effective passage retrieval model.",
"Jina-ColBERT is a ColBERT-style model but based on JinaBERT so it can support both 8k context length.",
"JinaBERT is a BERT architecture that supports the symmetric bidirectional variant of ALiBi to allow longer sequence length.",
"Jina-ColBERT model is trained on MSMARCO passage ranking dataset, following a very similar training procedure with ColBERTv2.",
"Jina-ColBERT achieves the competitive retrieval performance with ColBERTv2.",
"Jina is an easier way to build neural search systems.",
"You can use Jina-ColBERT to build neural search systems with ease.",
# Add more documents here to ensure the clustering work correctly
]
indexer.index(name=index_name, collection=documents)
tag在 RAGatouille 中使用 jina-colbert-v1-en
RAGatouille 是一个新的 Python 库,它简化了在 RAG 流程中使用高级检索方法的过程。它的设计注重模块化和易于集成,使用户能够无缝使用最先进的研究成果。RAGatouille 的主要目标是简化在 RAG 流程中应用 ColBERT 等复杂模型的过程,让开发者无需深入了解底层研究就能使用这些方法。感谢 Benjamin Clavié,现在你可以轻松使用 jina-colbert-v1-en:
from ragatouille import RAGPretrainedModel
# Get your model & collection of big documents ready
RAG = RAGPretrainedModel.from_pretrained("jinaai/jina-colbert-v1-en")
my_documents = [
"very long document1",
"very long document2",
# ... more documents
]
# And create an index with them at full length!
RAG.index(collection=my_documents,
index_name="the_biggest_index",
max_document_length=8190,)
# or encode them in-memory with no truncation, up to your model's max length
RAG.encode(my_documents)
要了解更多关于 Jina-ColBERT 的详细信息,你可以访问其 Hugging Face 页面。
tag结论
ColBERT 在信息检索领域代表了一个重要的进步。通过 Jina-ColBERT 支持更长的上下文长度,同时保持与 ColBERT 后期交互方法的兼容性,它为希望实现最先进搜索功能的开发者提供了一个强大的选择。
结合 RAGatouille 库,它简化了将复杂检索模型集成到 RAG 流程中的过程,开发者现在可以轻松利用高级检索的能力,简化他们的工作流程并增强他们的应用。Jina-ColBERT 和 RAGatouille 的协同展示了在使高级 AI 搜索模型变得易于访问和高效使用方面的显著进展。