ColBERT和ColBERTv2:兼具Bi-encoder和cross-encoder优势的多向量排序模型

简介

ColBERT是一种多向量排序模型，因为引入了延迟交互机制(late interaction architecture)相比与cross-encoder效率提升了很多。ColBERTv2针对ColBERT的缺点进一步优化了性能和效率。在RAG大热的这一年，ColBERT也引起了一些关注。

ColBERT出自2020年的论文《ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT》

ColBERTv2出自2022年初的论文《ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction》

ColBERT

ColBERT是a ranking model based on contextualized late interaction over BERT的简称。

在介绍ColBERT的实现之前，先来理解一下延迟交互(late interaction)是什么意思，设有query q和document d，在延迟交互机制下，q和d分别被编码成两个上下文embedding集合，q和d的相关性在两个集合中高效地完成。下图对比了延迟交互机制与其他几种文本相似度匹配方法的区别。

• Figure 2(a) 是双塔结构，被原作者归类为representation-focused rankers，分别计算q和d的embedding向量之后，用向量相似度分数来估计q和d的相关性。
• Figure 2(b) 被原作者归类为interaction-focused rankers，该框架尝试建模q和d在词级别和短语级别的关系并通过DNN或者kernel的方式来进行匹配。
• Figure 2© 是Figure 2(b)的进阶版，典型结构如BERT架构，同时对 q 和 d 内以及 q 和 d 之间的单词之间的交互进行建模。
• Figure 2(d)是Colbert提出的延迟交互机制，它结合了前几项优点，Figure 2(a)代表的方法可以进行离线计算，而Figure 2(b)和Figure 2©代表的交互方法通常效果更好，ColBERT将两者结合起来提出了延迟交互的思路（详细思路见后文）。

ColBert 原理

ColBERT的整体架构图如上图(即上上图里的Figure2(d))，由一个query encoder $f_Q$、一个document encoder $f_D$和延迟交互机制构成。

ColBERT使用相同的BERT-based encoder将每一个query或document编码成多个embedding（a bag of embeddings），但是会对query和document作不同的处理。

1. query encoder

• 对于一个query q，使用BERT-based WordPiece分词器将其分词成$q_1{q}_2...q_l$，在BERT序列起始特殊token [CLS]后添加一个token [Q]。预定义一个query的token长度须为$N_q$（作者实验里取值为32） ，如果一个query的token数小于$N_q$，则用[mask] token将其填充为$N_q$个token；相反地，如果一个query的token个数大于$N_q$，则取其前$N_q$​个token。
• 将token序列输入到BERT，为每个token生成一个上下文表征。
• 将BERT输出的每一个token上下文表征经过一个无激活函数的线性层，生成m维embedding，m是一个比BERT输出隐藏层维度小很多的值（作者实验里取值为128）。
• 将每个token的m维输出embedding归一化到$L_2$ 范数为1。

2. document encoder

• 对于一个document d，使用BERT-based WordPiece分词器将其分词成$d_1{d}_2...d_n$，在BERT序列起始特殊token [CLS]后添加一个token **[D]**表明这是一个文档序列。
• 将文档token序列输入到BERT，为每个token生成一个上下文表征
• 将BERT输出的每一个token上下文表征经过一个无激活函数的线性层，生成m维embedding，m是一个比BERT输出隐藏层维度小很多的值（作者实验里取值为128）。
• 将每个token的m维输出embedding归一化到$L_2$ 范数为1。
• 使用预先定义的标点符号列表过滤掉属于标点符号的embedding。

所以给定$q=q_1{q}_2...q_l$和$d=d_1{d}_2...d_n$，生成embedding集 $E_q$和$E_d$的过程可用公式表示成下列式子（#是指[mask] token）:
$$\begin{array}{rl}& E_q:=\text{ Normalize( CNN( BERT(“[Q] }{q}_1{q}_2\dots q_l\#\#\dots \#\text{”})))(1)\\ & E_d:=\text{ Filter( Normalize( CNN( BERT(“[D] }\left[d_1{d}_2\dots d_n\text{”}\right)))(2)\end{array}$$

3. 延迟交互(late interaction)

在编码得到q和d的embedding集$E_q$和$E_d$（都是a bag of embeddings）后，q和d的相关分数$S_{q,d}$通过延迟交互计算，对每一个向量$v\in E_q$计算其与$E_d$中各向量的最大cosine相似度(MaxSim操作），再将这些相似度求和得到相关分数$S_{q,d}$。因为向量已经归一化了，直接进行dot-product就可以了，公式表示如下：
$$S_{q,d}:=\sum _{i\in \left[|E_q|\right]}\underset{j\in \left[|E_d|\right]}{\max }{E}_{q_i}\cdot E_{d_j}^T(3)$$

ColBERT如何训练

ColBERT的延迟交互机制没有可训练参数，训练时对BERT进行微调，并从头训练线性层参数以及[Q]和[D]的embedding。

给定三元组$<q,d^{+},d^{-}>$ ，代表query q， 正例文档$d^{+}$、负例文档$d^{-}$，ColBERT为两个文档$d^{+}$和$d^{-}$分别生成一个分数，并使用交叉熵来优化（论文中使用的描述是pairwise softmax cross-entropy loss，实现上使用torch.nn.CrossEntropyLoss()，所有label取值都是正例的位置0)

ColBERT 如何使用

离线索引

ColBERT在设计时考虑将query和document的计算尽可能分离，所以可以先计算document embedding并存储。

使用document encoder $f_D$将文档编码成embedding并以32-bit或16-bit来存储，过程中使用了如下优化技巧来提升索引速度。

• 使用多个GPU并行编码，在batching时，将所有文档pad到一个batch内最长文档。在预处理时，会将文档按照长度排序，每次将长度相差不多的文档放入到一个batch。
• 将BERT的WordPiece分词预处理操作在CPU上并行处理。

用ColBERT 实现top-k Re-ranking

假设对于给定query q，我们有k个文档需要排序。因为k相对较小(比如k=1000)，主要通过GPU来按照公式3来暴力计算这k个文档相对于query的分数。

cross-encoder需要经过k次长度$l=|q|+|d_i|$的BERT编码，而ColBERT只需要一次长度短的多的$l=|q|$的BERT编码，相比之下效率提升了很多（下图是论文中的性能比较示意）。

用ColBERT 实现top-k 端到端的检索

假设需要从N(如N=10,000,000)个文档集合中检索出top-k个文档，k是比N小得多的数值。因为文档集合非常大，此时就没有办法直接暴力匹配来计算了。解决办法是使用向量相似度检索如faiss。

• 首先将文档embedding 使用faiss的IVFPQ来索引，这里要保留每一个embedding与原文档之间的映射关系。
• 使用两阶段过程来检索top-k个文档。在第一阶段，并行地进行$N_q$个向量相似度查询(即对$E_q$里的$N_q$个向量分别进行查询)，检索出$\text{top}-k^{\prime }$个文档向量（$k^{\prime }=k/2$)。将这些文档向量映射到其文档id，将得到$N_q\times k^{\prime }$个文档id，容易想到这里面会有$K\le N_q\times k^{\prime }$个文档是唯一的，K个文档会包括一个或多个embedding与query embedding非常相似。
• 在第二阶段，就可以在K个文档上进行如前一节re-ranking一样的暴力计算排序了。

ColBERTv2

ColBERT相比与cross-encoder提升了速度，但是其多向量会占用很多空间，所以ColBERTv2采用残差压缩机制(residual compression mechanism）对ColBERT进行改进在保留性能同时减少了6-10倍的空间占用。

ColBERTv2原理

ColBERTv2的整体架构与ColBERT一模一样，同样由一个query encoder $f_Q$、一个document encoder $f_D$​和延迟交互机制构成。具体实现方式在此不赘述，参见前文。

Supervision

ColBERT是使用三元组$<q,d^{+},d^{-}>$来训练优化的，在之后有很多相关研究表明采用hard negative等方法可以提高检索性能，所以ColBERTv2在模型训练上采用如下优化：

• 先用ColBERT里的三元组训练一个ColBERT模型，并将训练集里的passage编码后使用后文会介绍的ColBERTv2压缩方法来索引。
• 对于每一个训练集query，检索出top-k passages。将这些query-passage对输出一个cross-encoder rerank来打分，作者使用的cross-encoder rerank是22M参数的MiniLM。
• 基于cross-encoder计算出的分数得到w-way tuples，w-way tuples包括一个query,一个高排名的passage（或者标注正样本），一个或多个低排名的passage。论文中作者取了w=64。
• 在接下来的ColBERT进一步训练中，使用KL散度损失将cross-encoder计算的分数蒸馏到ColBERT架构（GitHub code）。除了KL散度损失之外，还在每个GPU中使用in-batch negatives。

Representation

作者先假设ColBERT向量可聚类到捕捉token特定语义的区域。给定一个簇中心集合$C$， ColBERTv2将每个向量v索引编码成它最近的簇中心$C_t$和一个量化向量$\tilde{r}$（$\tilde{r}$是用来近似残差$r=v-C_t$的），在搜索时，使用簇索引$t$和$\tilde{r}$来恢复一个近似向量$\tilde{v}=C_t+\tilde{r}$。

在编码$\tilde{r}$时，将r的每个维度量化到1 bit 或2 bits。所以理论上，n维向量编码成b-bit后每个向量需要$\lceil \log |C|\rceil +bn$ bits大小的空间占用。在论文实验中，向量维度n=128，使用4 bytes 就可最多表示$2^{32}$个簇中心 ，再用16 bytes或32 bytes空间(对应b=1或b=2)来编码残差，所以每个向量的总占用空间为20 bytes或36 bytes。 而ColBERT使用16-bit精度表示128维向量需要256 bytes空间，所以ColBERTv2的压缩策略可以显著降低存储成本。

Indexing

索引阶段提前计算所有passage的embedding，经过一定处理以支持快速的最近邻搜索，ColBERTv2将索引分为如下三个阶段：

1. 簇中心选择(Centroid Selection)，簇中心集合的大小$|C|$是与$16\times \sqrt{n_{embeddings}}$最接近的2次幂，即其与语料集中所有的embedding个数$n_{embeddings}$成比例。从所有passage embedding中抽样一部分来进行k-means聚类得到$|C|$个簇中心。（在代码中抽样的passage个数为$16\times \sqrt{120\ast num\_passages}$））

2. Passage编码(Passage Encoding)，将语料库中所有的passage经过BERT encoder生成embedding集之后，使用前一节中的方法进行压缩表示并存储到磁盘。

3. Index Inversion，为了支持快速近邻向量检索，将属于同一个簇中心的embedding ID分组，即建立倒排列表保存到磁盘。这样有助于在查询时快速找到相似embedding。

Retrieval

Retrieval分为两步：

1. 在检索时，先得到一些候选passage，对于query的每个向量$Q_i$，先找到最近的$n_{probe}\ge 1$个簇中心。使用倒排索引得到与这些簇中心相似的passage embeddings，将它们解压缩并计算其与每一个query向量的余弦相似度。对每一个query向量将这些分数按照passage ID分组，然后将相同passage的分数进行max-reduce得到唯一分数（这个操作相当于ColBERT里的MaxSim的近似，是真实MaxSim的一个下界）。将所有query向量对应一个passage的分数求和即得到passage的分数，以此分数排序后得到top $n_{candidate}$个候选passage。

2. 将候选集中passage的embedding集解压缩还原之后，按照ColBERT 里介绍过的top-k reranking方法计算分数并排序返回结果。

总结

文本介绍了多向量排序模型ColBERT和ColBERTv2 ，ColBERT提出了延迟交互机制，ColBERTv2在ColBERT基础上使用更先进的训练方法来微调模型，并通过残差压缩方法大幅减少存储成本。 BGE-M3的多向量(multi-vector)能力与ColBERT的思路是一样的，看完本文相信对BGE-M3也更理解了。

注1：在ColBERT issue 有讨论说RoBERTa模型因为训练时没有下句预测任务，使用ColBERT的效果并不好

注2：在ColBERTv2论文中有试验表明ColBERT因为是多向量，所以在out of domain上的效果更好。

参考资料

1. Khattab, Omar, and Matei Zaharia. 2020. “ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT.” arXiv: Information Retrieval,arXiv: Information Retrieval, April.

2. Santhanam, Keshav, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. 2022. “ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction.” In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. doi:10.18653/v1/2022.naacl-main.272.

3. https://jina.ai/news/what-is-colbert-and-late-interaction-and-why-they-matter-in-search/

4. colbert github

5. ragatouille 也实现了colbert, llama-index和langchain等都有集成ragatouille

6. 知乎文章