RAG工具箱：ColBERT

前一段时间跑比赛时，苦于上分陷入瓶颈，于是便祭出了ML经典ensemble大法，其中一个分支便是选择的ColBERT，没想到，竟成了破局的关键；也由此补齐了我对“迟交互”类模型的认知。

在典型的检索系统中，我们常用的架构是 Retrieval + Reranking两阶段:

检索阶段使用轻量级的方法（如BM25或dense vector retriever）从海量语料中快速筛选出几百个候选文档，优先考虑速度和召回率；排序阶段则使用cross-encoder基于 query-document 全量交互，重新给候选集打分，得到精准的最终排序。

这么做虽运行良好，但存在瓶颈：第一阶段检索质量的限制意味着大量低质量候选被送入昂贵的第二阶段，而cross-encoder在大规模候选集上的计算成本奇高，高相关的文档如果没能被第一阶段召回，就无法被 reranker “捞回来”。

而multi-vector模型正好填补了这个检索流水线中的“空白中间层”——比 embedding 更精准，比 cross-encoder 更高效。相比单向量 embedding 模型将整个chunk压缩为一个向量，multi-vector模型保留了 **token 级别的表示，**可以捕捉到局部与 query 高度匹配的片段，也是能缓解“语义鸿沟”的原因。

而ColBERT 便是经典的多向量模型/late-interaction代表。

（可能在本文中会将ColBERT/late-interaction/multi-vector-model 混为一谈，毕竟只是实例-机制-类别的关系）

不同的“交互”层级

（借用ColBERT论文中的插图）

想要理解late-interaction，需要先理解过去的不同“交互”形式；

Representation-based Similarity （no-interaction，embedding）

也就是双塔， query和document独立计算出一个向量，比较两个向量的相似度获得分数S；在生成各自的向量前， Query 和 Document 之间没有任何信息交互，所有复杂的语义都被压缩进了一个单一的向量中（也就存在信息压缩的损失和瓶颈）

优点速度极快，但缺点也很明显，精度有限，所有信息被压缩到一个向量中。

All-to-all Interaction （cross-encoder）

目前主流基于BERT 实现的Reranker / Cross-Encoder

将查询和文档拼接成 [CLS] query_tokens [SEP] document_tokens [SEP]

输入到 BERT 模型中，[CLS] 特殊标记最终输出向量，将向量输入一个linear layer中，从而获得获得一个0～1间的标量值。

因为是在入口处做的交互，于是也被称为“Early, Deep, All-to-all Interaction”

目前而言效果比较好的排序方案，代价是计算速度慢

Late Interaction （ColBERT）

查询和文档是独立、并行地通过 BERT 编码器进行编码，生成各自的token级别向量表示

交互发生在编码之后。它拿着编码好的 Query 词元向量集合，去和编码好的 Document 词元向量集合进行轻量级的、后置的相似度计算（MaxSim 操作）

具体而言，独立编码阶段，Query $q$ 编码为一组向量 $E_q = {q_1, q_2, …, q_N}$, Document $d$ 编码为一组向量 $E_d = {d_1, d_2, …, d_M}$；

计算MaxSim阶段，对于 Query 中的每一个token向量$q_i$, 计算它与 Document 中所有token向量 $d_j$ 的余弦相似度

通过

$$ MaxSim(q_i, E_d) = max(CosSim(q_i, d_1), CosSim(q_i, d_2), …, CosSim(q_i, d_M)) $$

找到Query 中的 $q_i$ 这个token在 Document 中能找到的最大值 （即最佳匹配）

将 Query 中所有token的MaxSim相加，得到最终的 Query-Document 相关性总分

$$ Score(q, d) = \sum_{for \ i=1 \ to \ N} MaxSim(q_i, E_d) $$

在训练过程中，与 BERT Reranker 类似，ColBERT也使用三元组 (query, positive_doc, negative_doc) 进行对比学习。

ColBERT 使用

这里利用RAGatouille实现ColBERT的使用。RAGatouille集成了不同的ColBERT类模型，能够在现有的RAG-pipeline中开箱即用。

# 依赖
pip install ragatouille

不过由于colbert使用了低版本的transofmers，直接使用会报错：ImportError: cannot import name 'AdamW' from 'transformers'

建议直接修改源码~\site-packages\colbert\training\training.py

from

from transformers import AdamW, get_linear_schedule_with_warmup
optimizer = AdamW(filter(lambda p: p.requires_grad, colbert.parameters()), lr=config.lr, eps=1e-8)

from transformers.optimization import get_linear_schedule_with_warmup
optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, colbert.parameters()), lr=config.lr, eps=1e-8)

即可使用；

这里使用了官方提供的示例

from ragatouille import RAGPretrainedModel

RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")

from ragatouille.utils import get_wikipedia_page

my_documents = [get_wikipedia_page("Hayao_Miyazaki"), get_wikipedia_page("Studio_Ghibli")]
index_path = RAG.index(index_name="my_index", collection=my_documents)


"""
[Jun 30, 06:30:31] #> Creating directory .ragatouille/colbert/indexes/my_index 
[Jun 30, 06:30:33] [0] 		 #> Encoding 117 passages..
...
[Jun 30, 06:30:34] [0] 		 avg_doclen_est = 193.56410217285156 	 len(local_sample) = 117
[Jun 30, 06:30:34] [0] 		 Creating 2,048 partitions.
[Jun 30, 06:30:34] [0] 		 *Estimated* 22,646 embeddings.
[Jun 30, 06:30:34] [0] 		 #> Saving the indexing plan to .ragatouille/colbert/indexes/my_index/plan.json ..
used 19 iterations (0.4336s) to cluster 21515 items into 2048 clusters
...

100%|██████████| 1/1 [00:00<00:00, 1034.86it/s][Jun 30, 06:32:31] #> Optimizing IVF to store map from centroids to list of pids..
[Jun 30, 06:32:31] #> Building the emb2pid mapping..
[Jun 30, 06:32:31] len(emb2pid) = 22647

100%|██████████| 2048/2048 [00:00<00:00, 74322.18it/s][Jun 30, 06:32:31] #> Saved optimized IVF to .ragatouille/colbert/indexes/my_index/ivf.pid.pt
Done indexing!
"""

(on T4, 约2min完成对117个文档的index构建)

query：

query = "照樹 務"
RAG = RAGPretrainedModel.from_index(".ragatouille/colbert/indexes/my_index")
results = RAG.search(query)
results

"""
[{'content': 'Wonderbird. Visually, he was inspired by Kagoshima Publishing\'s Italian Mountain Cities and the Tiber Estuary, reflecting his love for Europe. Production ran for four months and the film was released on December 15, 1979; Miyazaki wished he could have had another month of production. It was well received; Animage readers voted it the best animation of all time—it remained in the top ten for more than fifteen years—and Clarisse the best heroine. In 2005, former princess Sayako Kuroda\'s wedding dress was reportedly inspired by Clarisse\'s, having been a fan of Miyazaki and his work. Several Japanese and American filmmakers were inspired by the film, prompting homages in other works.\nMiyazaki became a chief animation instructor for new employees at Telecom Animation Film, a subsidiary of Tokyo Movie Shinsha. and subsequently directed two episodes of Lupin the Third Part II under the pseudonym Teruki Tsutomu (照樹 務), which can read as "employee of Telecom". In his role at Telecom, Miyazaki helped train the second wave of employees.',
  'score': 15.0859375,
  'rank': 1,
  'document_id': '953a72c5-2adc-4e34-ab8a-ced047ee3140',
  'passage_id': 20},
...

]

"""

ColBERT 新作

自ColBERT在2020年提出以来，在“迟交互”方向已发展出多个变体和优化版本。

ColBERTv2 在原来的ColBERT基础上，增加了1）去噪监督策略，以提升模型的检索精度；2）残差压缩机制，大幅减小索引的存储体积

PLAID为降低在大规模语料下的检索延迟，引入 centroid interaction（把每个段落视为一个centroid ID 的bag of centroids，只用 query-to-centroid 分数在 token 级估算段落得分，从而在不解压残差的情况下快速筛掉低分段落）和centroid pruning（基于查询对所有 centroids 的分数，先剪掉与查询无关或分数很低的 centroid，从而稀疏化袋表示，进一步减少计算）

jinaai/jina-colbert-v2 采用Matryoshka Representation Learning，以支持不同输出向量大小（128、96 和 64）；同时支持8192 tokens的上下文长度。

MUVERA Weaviate为解决ColBERT中存在的巨大的内存占用和速度慢的问题，MUVERA 的目标是将多向量嵌入 $D$ 和查询 $Q$ 分别编码成单个向量 $d_{single}$ 和 $q_{single}$，并使得它们的点积结果 $d_{single} \cdot q_{single}$ 能够很好地近似原始的 $maxSim(D,Q)$ 相似度得分,将复杂的“多对多”向量搜索问题就简化为了“一对一”单向量搜索问题

ConstBERT 与ColBERT为每个词元存储一个向量不同，ConstBERT通过一个可学习的投影层，将一个文档的所有词元向量压缩并投影成一个固定数量的向量（32或64），而与文档的实际长度无关，这些新生成的固定数量的向量不再与原文中的特定词元直接对应，而是文档不同“语义方面”的抽象表示