ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT
Omar Khattab, Matei Zaharia
cs.IR, cs.CL
2020-04-27
查询和文档各自独立过BERT编码，再用MaxSim做延迟交互。MS MARCO段落检索上MRR@10为34.9，相对BERT-base重排快约170倍、每查询FLOPs少约一万四千倍。
2019年前后，把查询和文档拼在一起送进BERT，再用[CLS]打分，把MS MARCO的MRR@10从非BERT方法的高二十多推到三十五左右。代价是每个查询-文档对都要过一遍深度Transformer。Nogueira和Cho的BERT-base重排官方前1000，单条查询约10,700毫秒、97T FLOPs；BERT-large到32,900毫秒。检索多100毫秒就会伤体验和收入，这个成本在线上几乎用不起。
另一头，单向量表示可以把文档离线算好，但细粒度匹配没了，效果上不去。当时的选择像是：要么贵且准，要么便宜且粗。
ColBERT的做法是late interaction。查询和文档各自用BERT编成一组上下文向量，交互留到最后，而且要便宜、还要能剪枝。
两边共享一个BERT，输入前加[Q]或[D]。查询不足Nq=32个token就用[MASK]补齐，叫做query augmentation，让模型学着在mask位置上写出可匹配的扩展。文档不过mask，并丢掉标点向量。BERT输出再过一层无激活的线性层，把维度收到m=128，然后L2归一化，点积就是余弦。
打分是每个查询向量对文档所有向量取最大相似度（MaxSim），再对查询项求和。没有可训练的交互参数。训练用三元组，对正负文档的分数做pairwise softmax交叉熵。
因为文档编码与查询无关，9M篇MS MARCO段落可以离线建好。重排时查询只过一次BERT，再和已缓存的文档矩阵做批MaxSim。端到端检索则把所有文档向量丢进faiss的IVFPQ：每个查询向量先取回一批近邻文档，去重后再对这小集合做精确MaxSim。论文里端到端用平方L2。
重排官方BM25 top-1000，Dev集MRR@10：
| 方法 | MRR@10 Dev | 延迟 | FLOPs/query |
| KNRM | 19.8 | 3 ms | 592M |
| Duet | 24.3 | 22 ms | 159B |
| fastText+ConvKNRM | 29.0 | 27.7 Eval | 78B |
| BERT-base | 34.7 | 10,700 ms | 97T |
| BERT-base（同损失重训） | 36.0 | 10,700 ms | 97T |
| BERT-large | 36.5（Eval 35.9） | 32,900 ms | 340T |
| ColBERT | 34.9（Eval 34.9） | 61 ms | 7B |
相对BERT-base，延迟约1/170，FLOPs约1/13900。61毫秒里查询编码加交互只有13毫秒，剩下是把缓存向量搬到GPU。k越大，BERT越吃亏：k=10时约180倍FLOPs，k=1000时13,900倍，k=2000时23,000倍。
端到端从8.8M篇里直接取top-1000：ColBERT MRR@10 36.0，超过重排的34.8，因为召回更高。Recall@1000为96.8，对Anserini BM25的85.7、docTTTTTquery的94.7。延迟458毫秒。TREC CAR的MAP：BM25 15.3，BERT-base 31.0，BERT-large 33.5，ColBERT 31.3。
空间可以压。128维4字节重排占286 GiB、MRR 34.9；24维2字节只要27 GiB，MRR 33.9，掉1个点。四卡给MS MARCO建索引大约三小时。消融里，单向量[CLS]点积、把MaxSim换成平均相似度、去掉query augmentation，都明显低于完整ColBERT。
这是后来多向量检索一条线的起点。工程含义很直接：文档离线过一遍BERT，线上用向量索引加一层便宜的MaxSim，就能在接近交叉编码器的质量下把延迟从十几秒打到几十毫秒。今天的token-level检索、PLAID一类压缩，都还在这条「迟交互、可剪枝」的路上。
它不是全面超过BERT-large。Eval上34.9对35.9，差一个点。换来的是四个数量级的计算下降，以及可以不经过BM25直接从全库取回。
2019到2020年的实验，编码器是BERT-base，没有指令模型、没有长文档、没有多语言。端到端458毫秒仍明显高于BM25的62毫秒，faiss还用了全部CPU核。索引体积即使压到27 GiB也远大于倒排表。MaxSim对每个查询项独立找最像的文档token，解释性好，但论文没有证明这就是相关性的正确分解。TREC CAR为了避免维基泄漏用了另训的BERT-large初始化，和MS MARCO设定不完全可比。后续工作已经把存储和延迟再压过一轮，这篇应读成范式，而不是今天的SOTA数字。