ColBERT: Eficient and Efective Passage Search via Contextualized Late Interaction over BERT
1.介绍
预训练的LM在IR上计算成本太高。本文考虑降低计算成本的同时,保证检索性能,并且结合基于表示和交互的模型(如下图
d
d
d)。
预训练LM带来的显著提升是以计算成本为代价的:下图展示了各个检索模型的计算成本,所有的model在Tesla V100 GPU上对MS MARCO上Ranking
为了协调信息检索中的效率和语境化,作者提出了
C
C
C
o
o
o
l
l
l
B
B
B
E
E
E
R
R
R
T
T
T:基于上下文化的后期交互的排序模型–contextualized late interaction over BERT 。ColBERT提出了一种新的后期交互范式,用于估计查询
q
q
q和文档
d
d
d之间的相关性。在后期交互下,
q
q
q和
d
d
d被分别编码到两组上下文嵌入中,并且相关性是使用两组之间廉价且
p
p
p
r
r
r
u
u
u
n
n
n
i
i
i
n
n
n
g
g
g-
f
f
f
r
r
r
i
i
i
e
e
e
n
n
n
d
d
d
l
l
l
y
y
y的计算来评估的——也就是说,快速计算能够在不穷举评估每个可能的候选项的情况下进行排名。整体框架看下图右👉IR模型对比
IR模型对比:
(
d
d
d):每个查询 embedding 都通过MaxSim操作符与所有文档嵌入交互,MaxSim操作符计算最大相似度(例如cos),并且这些操作符的标量输出在查询项之间求和。这个范例允许
C
C
C
o
o
o
l
l
l
B
B
B
E
E
E
R
R
R
T
T
T可以利用基于LM的深度表示,同时将编码文档的成本转为offline,并在所有已排序文档中一次性摊销编码查询的成本。 此外,它使Colbert能够利用向量相似度搜索索引(FAISS)直接从大型文档集中检索前
k
k
k个结果,从而大大改善了仅基于术语的re-rank模型的召回率 。
2. Architecture
一个query encoder
f
Q
f_Q
fQ,编码query
q
q
q到
E
q
E_q
Eq。
一个document encoder
f
D
f_D
fD,编码文档
d
d
d到
E
d
E_d
Ed。
外加一个late 交互机制计算
q
q
q和
d
d
d之间的相关性。
3.2 Query & Document Encoders
query encoder
BERT的encoder,在query和document前分别加入[Q]和[D]token。
给定一个
q
q
q,bert tokenizer :
q
1
q_1
q1,…,
q
l
q_l
ql,预处理:[CLS] [Q] [
q
1
q_1
q1],…,[
q
l
q_l
ql]
如果
q
q
q的长度小于max_len,则[mask] token补齐。否则裁剪到max_len长度的
q
q
q。这一步称为查询增强。
每个token经过BERT后,将输出的上下文表示通过一个没有激活函数的线性层传递上下文化的输出表示。该层用于控制 C C C o o o l l l B B B E E E R R R T T T嵌入的维度,输出大小 m m m维的嵌入。 m m m是远远小于bert隐藏层维度的。
这一步主要为了控制文档和查询的嵌入的占用空间和执行效率,虽然这样做会在一定程度上影响嵌入的质量。
Document Encoder.
同是BERT的encoder,与query encoder不同之处在于不额外的使用[mask] token。并且在通过BERT时过滤掉文档中的标点符号的嵌入(作者认为标点符号的嵌入对于检索的有效性是不必要的,这里的操作是根据一个预先定义的标点符号列表进行过滤)
这里有待商讨?
总结:给定
q
q
q =
q
0
q_0
q0,
q
1
q_1
q1,…,
q
l
q_l
ql
d
d
d =
d
0
d_0
d0,
d
1
d_1
d1,…,
d
n
d_n
dn
首先用BERT encoder 计算
E
q
E_q
Eq,
E
d
E_d
Ed ,(#表示[mask] token):
3.3 Late Interaction
maximum similarity 实现:
C C C o o o l l l B B B E E E R R R T T T是可区分的端到端的。通过对BERT编码器进行微调,并使用Adam 优化器从头开始训练附加参数(即线性层和[Q]和[D]标记的嵌入)。
注意,交互机制没有可训练的参数。
给定带有查询 q q q、正文档 d + d^+ d+和负文档 d − d^- d−的三元组< q q q、 d + d^+ d+、 d − d^- d−>, C C C o o o l l l B B B E E E R R R T T T用于为每个文档单独生成一个分数,并通过 d + d^+ d+和 d − d^- d−的计算分数上的成对sofmax交叉熵损失进行优化。
3.4 Ofline Indexing: Computing & Storing Document Embeddings
首先,利用多个GPU并行对成批量文档进行更快的编码。批量处理时,将所有文档填充到批内文档的最大长度。为了更有效地限制每批的序列长度,索引器将文档分组(例如,B = 100,000)进行处理。它按长度对这些文档进行排序,然后通过编码器提供具有相同长度的成批文档(例如,b = 128)。
Pytoch 中的BucketIterator
一旦产生了文档表示,就使用32位或16位值来表示每个维度,并将它们保存到磁盘上。这些表示要么简单地从磁盘加载以进行排序,要么随后分别被索引以进行向量相似性搜索。
3.5 Top-k Re-ranking with ColBERT
k=1000
首先,将索引文档表示加载到内存中,将每个文档表示为嵌入矩阵。
给定一个查询
q
q
q,计算它的上下文化嵌入
E
q
E_q
Eq,同时,将文档表示集合到由
k
k
k个文档矩阵组成的三维tensor
D
D
D 中([k,seq_len,m-dim])。然后将
k
k
k个文档填充到它们的最大长度,以便于批量操作,并将张量
D
D
D移动到GPU的内存中。在GPU上,计算一个batch 内的
E
q
E_q
Eq和
D
D
D的点积。输出具体化为一个三维张量,它是
q
q
q和每个文档之间的交叉匹配矩阵的集合。为了计算每个文档的分数,通过一个max pooling来减少其跨文档项的矩阵,并通过求和来减少跨查询项的矩阵。最后,我们按照总分对k个文档进行排序。
3.6 End-to-end Top-k Retrieval with ColBERT
这部分关注的情况是,要排序的文档数量太大(例如,
N
N
N = 10,000,000个文档,选择top-
k
k
k文档,
k
k
k
《
《
《
《
《
《
N
N
N),无法对每个可能的候选文档进行详尽的评估,尤其是当只对得分最高的文档感兴趣时。
这里直接调用强大的FAISS即可。
htps://github.com/facebookresearch/faiss
在3.4节索引结束时,我们维护从每个嵌入到其原始文档的映射,然后将所有文档嵌入索引到faiss中。进行大规模向量相似性搜索。
随后,在提供查询时,使用两阶段过程从整个集合中检索前k个文档。在第一阶段,同时向FAISS索引发出
N
q
N_q
Nq个向量相似性查询(对应于等式中的每个嵌入
E
q
E_q
Eq)。检索该向量的top-
k
’
k^’
k’(
k
’
k^’
k’=
k
k
k/2)匹配。然后产生
N
q
N_q
Nq×
k
’
k^’
k’标识符。其中仅仅有
K
K
K≤
N
q
N_q
Nq×
k
’
k^’
k’是唯一标识符,这些
K
K
K个文档可能包含一个或多个与查询嵌入高度相似的嵌入。对于第二阶段,通过以第3.5节中描述的通常方式仅对那些
K
K
K文档进行彻底的重新排序来完善此集合。
在基于faiss的实现中,使用了IVFPQ索引。这个索引基于
k
k
k均值聚类将嵌入空间划分为
P
P
P(例如,
P
P
P = 1000)个单元,然后基于所选的向量相似性度量将每个文档嵌入分配到其最近的单元。当搜索单个查询嵌入的top-
k
’
k^’
k’个匹配时,只搜索最近的
p
p
p(例如,
p
p
p = 10)个分区。为了提高存储效率,每个嵌入被分成
s
s
s(例如,
s
s
s = 16)个子向量,每个子向量用一个字节表示。此外,索引在这个压缩域中进行相似性计算,从而导致更便宜的计算和更快的搜索。
4 EXPERIMENTAL EVALUATION
- RQ1:在典型的重新排序设置中, C C C o o o l l l B B B E E E R R R T T T在多大程度上能够弥合高效神经模型(BERT)之间的现有差距?
- RQ2:除了重新排序之外, C C C o o o l l l B B B E E E R R R T T T能否有效地支持直接从一个大集合中进行端到端检索?
- RQ3: C C C o o o l l l B B B E E E R R R T T T的每个组件(例如,后期交互、查询增强)对其质量有什么贡献?
- RQ4:就内部计算和内存开销而言, C C C o o o l l l B B B E E E R R R$T的索引相关成本是多少?
4.1.1 Datasets & Metrics.
Datasets :MS MARCO Ranking 和TREC-CAR
Metrics:MRR@10
4.1.2 Implementation.
transformers库
微调BERT learning rate = 3×10
−
^-
−
6
^6
6 , batch_size = 32 , 固定每个查询embedding数量为
N
q
N_q
Nq = 32,
C
C
C
o
o
o
l
l
l
B
B
B
E
E
E
R
R
R
T
T
T隐藏层维度为128。
在re-ranking结果中,在嵌入中使用每个维度4个字节,并使用余弦作为向量相似性函数。
在end-to-end 模型中,使用L2距离,因为faiss指数在基于L2的检索中更快。对于faiss索引,将分区数设置为
P
P
P =2,000,并搜索每个查询嵌入的最近
p
p
p = 10,以检索每个查询嵌入的
k
’
k^’
k’ =
k
k
k = 1000个文档向量。将每个嵌入分成16个子向量,每个子向量用一个 字节编码。为了表示用于端到端检索过程第二阶段的索引,使用每个维度16-bit的值。
4.2 Qality–Cost Tradeof: Top-k Re-ranking
MS MARCO 的“re-rank”结果。每个神经模型重新排列BM25产生的前1000个结果。报告的延迟仅用于重新排序。
FLOPs:floating point operations,浮点运算数,计算量,衡量算法/模型的复杂度,理解为FLOPS在时间上的积分
表2:在MS MARCO上的端到端检索结果。每个模型直接从整个880万个文档集合中检索每个查询的前1000个文档。
总结
ColBERT采用情景化后期交互在深度PLMs(特别是BERT)进行高效检索。通过使用类BERT独立的对查询和文档进行编码,再使用一种廉价而且强大的交互步骤去建模它们的细粒度的相似性,可以利用深度PLMs的表现力而大大加快查询处理。此外,这样做允许使用ColBERT直接从大型文档集合进行端到端神经检索。与现有的基于bert的模型相比,ColBERT的速度要快170倍以上,需要的FLOPs/query要少14000倍,而且对质量的影响很小,而且比所有非bert的基线都要好。