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

最新推荐文章于 2024-10-14 08:20:52 发布

chencjiajy

最新推荐文章于 2024-10-14 08:20:52 发布

阅读量2.6k

点赞数 19

分类专栏：深度学习文章标签：论文阅读 RAG 排序模型 embedding

本文链接：https://blog.csdn.net/beingstrong/article/details/137988211

版权

深度学习专栏收录该内容

45 篇文章 3 订阅

订阅专栏

文章目录

简介

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作不同的处理。

query encoder

对于一个query q，使用BERT-based WordPiece分词器将其分词成 $q_1q_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。

document encoder

对于一个document d，使用BERT-based WordPiece分词器将其分词成 $d_1d_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_1q_2... q_l$ 和 $d=d_1d_2... d_n$ ，生成embedding集 $E_q$ 和 $E_d$ 的过程可用公式表示成下列式子（#是指[mask] token）:
$\begin{aligned} & \left.\left.\left.E_q:=\text { Normalize( CNN( BERT(“[Q] } q_1 q_2 \ldots q_l \# \# \ldots \# \text{”}\right)\right)\right) \qquad (1)\\ & \left.\left.E_d:=\text { Filter( Normalize( CNN( BERT(“[D] }\left[d_1 d_2 \ldots d_n \text{”}\right)\right)\right) \qquad (2) \end{aligned}$

延迟交互(late interaction)

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

簇中心选择(Centroid Selection)，簇中心集合的大小 $∣ C ∣$ 是与 $16 \times \sqrt{n_{embeddings}}$ 最接近的2次幂，即其与语料集中所有的embedding个数 $n_{embeddings}$ 成比例。从所有passage embedding中抽样一部分来进行k-means聚类得到 $∣ C ∣$ 个簇中心。（在代码中抽样的passage个数为 $16 \times \sqrt{120*num\_passages}$ ））
Passage编码(Passage Encoding)，将语料库中所有的passage经过BERT encoder生成embedding集之后，使用前一节中的方法进行压缩表示并存储到磁盘。
Index Inversion，为了支持快速近邻向量检索，将属于同一个簇中心的embedding ID分组，即建立倒排列表保存到磁盘。这样有助于在查询时快速找到相似embedding。

Retrieval

Retrieval分为两步：

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

参考资料

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.
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.
https://jina.ai/news/what-is-colbert-and-late-interaction-and-why-they-matter-in-search/
colbert github
ragatouille 也实现了colbert, llama-index和langchain等都有集成ragatouille
知乎文章