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 f_Q fQ、一个document encoder f D f_D fD和延迟交互机制构成。

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 q_1q_2... q_l q1q2...ql,在BERT序列起始特殊token [CLS]后添加一个token [Q]。预定义一个query的token长度须为 N q N_q Nq(作者实验里取值为32) ,如果一个query的token数小于 N q N_q Nq,则用[mask] token将其填充为 N q N_q Nq个token;相反地,如果一个query的token个数大于 N q N_q Nq,则取其前 N q N_q Nq​个token。
  • 将token序列输入到BERT,为每个token生成一个上下文表征。
  • 将BERT输出的每一个token上下文表征经过一个无激活函数的线性层,生成m维embedding,m是一个比BERT输出隐藏层维度小很多的值(作者实验里取值为128)。
  • 将每个token的m维输出embedding归一化到 L 2 L_2 L2 范数为1。
  1. document encoder
  • 对于一个document d,使用BERT-based WordPiece分词器将其分词成 d 1 d 2 . . . d n d_1d_2... d_n d1d2...dn,在BERT序列起始特殊token [CLS]后添加一个token **[D]**表明这是一个文档序列。
  • 将文档token序列输入到BERT,为每个token生成一个上下文表征
  • 将BERT输出的每一个token上下文表征经过一个无激活函数的线性层,生成m维embedding,m是一个比BERT输出隐藏层维度小很多的值(作者实验里取值为128)。
  • 将每个token的m维输出embedding归一化到 L 2 L_2 L2 范数为1。
  • 使用预先定义的标点符号列表过滤掉属于标点符号的embedding。

所以给定 q = q 1 q 2 . . . q l q=q_1q_2... q_l q=q1q2...ql d = d 1 d 2 . . . d n d=d_1d_2... d_n d=d1d2...dn,生成embedding集 E q E_q Eq E d E_d Ed的过程可用公式表示成下列式子(#是指[mask] token):
E q : =  Normalize( CNN( BERT(“[Q]  q 1 q 2 … q l # # … # ” ) ) ) ( 1 ) E d : =  Filter( Normalize( CNN( BERT(“[D]  [ d 1 d 2 … d n ” ) ) ) ( 2 ) \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} Eq:= Normalize( CNN( BERT(“[Q] q1q2ql###)))(1)Ed:= Filter( Normalize( CNN( BERT(“[D] [d1d2dn)))(2)

  1. 延迟交互(late interaction)

在编码得到q和d的embedding集 E q E_q Eq E d E_d Ed(都是a bag of embeddings)后,q和d的相关分数 S q , d S_{q,d} Sq,d通过延迟交互计算,对每一个向量 v ∈ E q v \in E_q vEq计算其与 E d E_d Ed中各向量的最大cosine相似度(MaxSim操作),再将这些相似度求和得到相关分数 S q , d S_{q,d} Sq,d。因为向量已经归一化了,直接进行dot-product就可以了,公式表示如下:
S q , d : = ∑ i ∈ [ ∣ E q ∣ ] max ⁡ j ∈ [ ∣ E d ∣ ] E q i ⋅ E d j T ( 3 ) 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) Sq,d:=i[Eq]j[Ed]maxEqiEdjT(3)

ColBERT如何训练

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

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

ColBERT 如何使用

离线索引

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

使用document encoder f D f_D fD将文档编码成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 ∣ l=|q| + |d_i| l=q+di的BERT编码,而ColBERT只需要一次长度短的多的 l = ∣ q ∣ l=|q| l=q的BERT编码,相比之下效率提升了很多(下图是论文中的性能比较示意)。

在这里插入图片描述

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

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

  • 首先将文档embedding 使用faiss的IVFPQ来索引,这里要保留每一个embedding与原文档之间的映射关系。
  • 使用两阶段过程来检索top-k个文档。在第一阶段,并行地进行 N q N_q Nq个向量相似度查询(即对 E q E_q Eq里的 N q N_q Nq个向量分别进行查询),检索出 top − k ′ \text{top}-k^{\prime} topk个文档向量( k ′ = k / 2 k^{\prime}=k/2 k=k/2)。将这些文档向量映射到其文档id,将得到 N q × k ′ N_q \times k^{\prime} Nq×k个文档id,容易想到这里面会有 K ≤ N q × k ′ K \leq N_q \times k^{\prime} KNq×k个文档是唯一的,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 f_Q fQ、一个document encoder f D f_D fD​和延迟交互机制构成。具体实现方式在此不赘述,参见前文。

Supervision

ColBERT是使用三元组 < q , d + , d − > <q, d^+, d^-> <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 C C, ColBERTv2将每个向量v索引编码成它最近的簇中心 C t C_t Ct和一个量化向量 r ~ \tilde{r} r~ r ~ \tilde{r} r~是用来近似残差 r = v − C t r=v-C_t r=vCt的),在搜索时,使用簇索引 t t t r ~ \tilde{r} r~来恢复一个近似向量 v ~ = C t + r ~ \tilde{v} = C_t + \tilde{r} v~=Ct+r~

在编码 r ~ \tilde{r} r~时,将r的每个维度量化到1 bit 或2 bits。所以理论上,n维向量编码成b-bit后每个向量需要 ⌈ log ⁡ ∣ C ∣ ⌉ + b n \lceil \log|C| \rceil + bn logC+bn bits大小的空间占用。在论文实验中,向量维度n=128,使用4 bytes 就可最多表示 2 32 2^{32} 232个簇中心 ,再用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 ∣ |C| C是与 16 × n e m b e d d i n g s 16 \times \sqrt{n_{embeddings}} 16×nembeddings 最接近的2次幂,即其与语料集中所有的embedding个数 n e m b e d d i n g s n_{embeddings} nembeddings成比例。从所有passage embedding中抽样一部分来进行k-means聚类得到 ∣ C ∣ |C| C个簇中心。(在代码中抽样的passage个数为 16 × 120 ∗ n u m _ p a s s a g e s 16 \times \sqrt{120*num\_passages} 16×120num_passages ))

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

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

Retrieval

Retrieval分为两步:

  1. 在检索时,先得到一些候选passage,对于query的每个向量 Q i Q_i Qi,先找到最近的 n p r o b e ≥ 1 n_{probe} \geq 1 nprobe1个簇中心。使用倒排索引得到与这些簇中心相似的passage embeddings,将它们解压缩并计算其与每一个query向量的余弦相似度。对每一个query向量将这些分数按照passage ID分组,然后将相同passage的分数进行max-reduce得到唯一分数(这个操作相当于ColBERT里的MaxSim的近似,是真实MaxSim的一个下界)。将所有query向量对应一个passage的分数求和即得到passage的分数,以此分数排序后得到top n c a n d i d a t e n_{candidate} ncandidate个候选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. 知乎文章

  • 14
    点赞
  • 12
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值