ICLR 2020 | 可提速3000倍的全新信息匹配架构（附代码复现）

2020 年 4 月 4 日 PaperWeekly

©PaperWeekly 原创 · 作者｜周树帆

学校｜上海交通大学硕士生

研究方向｜自然语言处理

今天聊一篇 FAIR 发表在 ICLR 2020 上的文章：Poly-encoders: Transformer Architectures and Pre-training Strategies for Fast and Accurate Multi-sentence Scoring。

论文标题：Poly-encoders: Transformer Architectures and Pre-training Strategies for Fast and Accurate Multi-sentence Scoring

论文来源：ICLR 2020

论文链接：https://arxiv.org/abs/1905.01969

和一些花里胡哨但是没有卵用的论文不同，这篇文章可谓大道至简。该文用一种非常简单但是有效的方式同时解决了 DSSM 式的 Bi-encoder 匹配质量低的问题和 ARC-II、BERT 等交互式的 Cross-encoder 匹配速度慢的问题。

背景

众所周知，常见的搜索、检索式问答、自然语言推断等任务，它们本质上都是一种相关性匹配任务：给定一段文本作为 query，然后匹配出最为相关的文档或答案然后返回给用户。

目前主流的文本相关性匹配架构有两大类：以 DSSM 为代表的 Siamese Network 架构、以及形如 ARC-II、ABCNN 或 BERT（基于 Self-Attention）的交互式匹配架构。

1.1 Siamese Network

如图 1 所示，Siamese Network 式（本篇文章又称其为 Bi-encoder）的匹配方案会利用 2 个网络分别将 query 和 candidates 编码成

和

，最后再通过一个相关性判别函数（通常为 cosine）计算两个 vec 之间的相似度。

这种方案的最大特点就是 query 和 candidates 直到最后的相关性判别函数时才发生交互，所以会对模型的匹配性能产生一定的影响。

但是这种完全独立的编码方式使得我们可以离线计算好所有 candidates 的向量，线上运行时只需计算 query 的向量然后匹配已有向量即可。总的来说，这种方案匹配速度极快，但是匹配质量不能达到最佳。

▲ 图1. Siamese Network（本篇论文又称其为Bi-encoder）

1.2 交互式匹配

如图 2 所示，交互式匹配（本文记作 Cross-encoder）的核心思想是则是 query 和 candidates 时时刻刻都应相互感知，相互交融，从而更深刻地感受到相互之间是否足够匹配。

早期的交互方案如 ARC-II、ABCNN 等会计算

和

之间的word embedding相似度、Q、C 分别过 RNN 之后的

、

之间的相似度，最后再用一些 CNN 之类的方法整合结果，然后用 MLP 做二分类判别是相关还是不相关。

▲ 图2. 交互式匹配示意图(图中为ARC-II)

另外在 BERT 兴起之后，如图 3 所示般将 query 和 candidate 拼成一句话，然后利用 self-attention 完成 query 和 candidate 之间的交互的模型也大量涌现，并且取得了非常显著的成果。本篇论文实现的 Cross-encoder 也是基于图 3 的架构。

相较于 Siamese Network，这类交互式匹配方案可以在 Q 和 C 之间实现更细粒度的匹配，所以通常可以取得更好的匹配效果。

但是很显然，这类方案无法离线计算 candidates的表征向量，每处理一个 query 都只能遍历所有 (query, candidate) 的 pairs 依次计算相关性，所以这类方案相当耗时（当然也有很多提速手段，不过那不是本文的重点）。

▲ 图3. Cross-encoder

Poly-Encoder

Bi-encoder (Siamese Network) 和 Cross-encoder（交互式网络）都有各自显著的优点和缺点，而本文提出的 Poly-encoder 架构同时集成了两类方案的优点并避免了缺点。

▲ 图4. Poly-encoder

Poly-encoder 如图 4 所示。Poly-encoder 的思想非常简单（简单到论文里仅用了 2 段文字），按我的个人理解描述：

Bi-encoder 的主要问题在于它要求 encoder 将 query 的所有信息都塞进一个固定的比较 general 的向量中，这导致最后

和

计算相似度时已经为时过晚，很多细粒度的信息丢失了（e.g. query 为“我要买苹果”），所以无法完成更精准的匹配。

这就有点像 word2vec 静态词向量：即使一个词有多种语义，它的所有语义也不得不塞进一个固定的词向量。

为了克服这个问题，Poly-encoder 的方案就是每个 query 产生 m 个不同的

，接着再根据

动态地将 m 个

集成为最终的

（其实有点像封面图那样，有一点用 m 个向量组合出最终的 Low Poly（baike.baidu.com/item/Lo）化向量的味道），最后再计算

和

的匹配度。

用论文里的话来说：

论文中的 ctxt 指代 context，相当于 query；cand 指代 candidate。上面这段论文建议我们可以随机初始化 m 个通过 dot product 计算 attention，从而将长度为 N 的 context 编码成 m 个向量

(即

)。

接着：

我们再用 candidate 对应的向量

计算 m 个

的 attention，进而得到最终的

。

很显然，Poly-encoder 架构在实际部署时是可以离线计算好所有 candidates 的向量的，所以只需要计算 query 对应的 m 个

向量，再通过简单的 dot product 就可以快速计算好对应每个 candidate 的“动态的”

向量。

看起来 Poly-encoder 享有 Bi-encoder 的速度，同时又有实现更精准匹配的潜力。我们通过实验来一探究竟。

实验

本文选择了检索式对话数据集 ConvAI2、DSTC 7、Ubuntu v2 数据集以及 Wikipedia IR 数据集进行实验。

训练 Bi-encoder 和 Poly-encoder 时由于这两类模型的特性，负采样方式为：在训练过程中，使用同一个 batch 中的其他 query 对应的 response 作为负样本（如果难以理解，可以稍后结合复现代码来理解）。

而 Cross-encoder 的负采样方式为：在开始训练之前，随机采样 15 个 responses 作为负样本。

3.1 检索质量

图5给出了一些 baseline 模型以及本文的基于预训练 BERT 的 Bi-encoder、Poly-encoder 以及 Cross-encoder 在各个数据集上的表现。

当然我们很容易发现，本文的所有模型由于以预训练的BERT为基础，他们的表现都要显著超出不使用 BERT的那些 baseline 们。所以我们只需要关注 Bi、Poly 和 Cross 三种架构之间的表现差异即可。

实验结果表明即使仅增设少数几个 code（用于计算 attention 产生向量），Poly-encoder 的表现也要远优于 Bi-encoder。

实验结果还表明，Poly-encoder 的表现会随着 code 个数的增加而逐渐增加，并且慢慢逼近 Cross-encoder 的结果（个人认为 Cross-encoder 的表现应该是 Poly-encoder 的上界，不过偶尔也可能会因为一些偶然因素导致 Poly-encoder 反超 Cross-encoder 的情况）。

另外，为了体现 Cross-encoder 在速度上的局限性，作者还很有意思地跳过了 Cross-encoder 在 Wikipedia IR 上的测评并写到：“In addition, Cross-encoders are also too slow to evaluate on the evaluation setup of that task, which has 10k candidates”。

▲ 图5. 模型表现汇总

3.2 检索速度

图 5 的实验结果已经表明 Poly-encoder 的检索质量明显优于 Bi-encoder 架构，且能逼近 Cross-encoder 架构的效果。剩下的关键问题就是 Poly-encoder 是否会显著增加检索耗时？

图 6 给出了各模型在 ConvAI2 数据集上的检索耗时。

Bi-encoder 理所当然是最快的架构，当 candidates 为 100k 时，在 CPU 和 GPU 环境下其检索耗时分别为 160ms 和 22ms；而 Cross-encoder 显然是最慢的一个：同样实验条件下其检索耗时分别约为 2.2M (220 万) ms 和 266K (26.6 万) ms。

反观 Poly-encoder，以 Poly-encoder 360 为例，该模型可以达到远超 Bi-encoder、接近甚至反超 Cross-encoder 的检索质量，但其检索速度确比 Cross-encoder 足足快了约 2600-3000 倍！

▲ 图6. 各模型在ConvAI2数据集上的检索耗时

论文小结

总的来说，本文的出发点就是希望找到一个速度快但质量不足的 Bi-encoder 架构和质量高但速度慢的 Cross-encoder 架构的折中。

本文提出的 Poly-encoder 的核心思想虽然非常简单，但是却十分有效（亲测），确实在很多场景下可以作为 Bi-encoder 的替代，甚至在一些对速度要求较高的场景下可以作为 Cross-encoder 的替代。

方案简洁固然是本文的一大优点，不过这也给未来的研究留下了空间。相信未来很快就会有许多基于 Poly-encoder 的改进版出现。

复现结果分享

在读完论文后的第一时间，我就尝试了复现工作。我的复现结果表明，Poly-encoder 不管是收敛速度还是模型上限，都要显著优于 Bi-encoder，且 Poly-encoder 几乎不增加额外的显存负担，对训练速度的影响也几乎可以忽略。完整代码位于：

https://github.com/sfzhou5678/PolyEncoder

5.1 关键代码分析

Poly-encoder 的实现非常简单，只需在 Bi-encoder 的基础上略加修改即可。接下来我将介绍实现 Poly-encoder 的核心代码。

我们首先用 nn.embedding 来作为 m 个 poly_codes 的值, 然后 forward 的时候根据m的值产生对应个数的 poly_codes，这些 codes 将用于计算不同的 attention weights，以产生多个 vec_ctxt（即 vec_q）。

这里我令 poly_code_ids+=1 是为了让 context_encoder 和 response_encoder 对称，所以把 0 号 id 留给了 response_encoder。

self.poly_code_embeddings = nn.Embedding(self.poly_m + 1, config.hidden_size)

poly_code_ids = torch.arange(self.poly_m, dtype=torch.long, device=context_input_ids.device)
poly_code_ids += 1
poly_code_ids = poly_code_ids.unsqueeze(0).expand(batch_size, self.poly_m)
poly_codes = self.poly_code_embeddings(poly_code_ids)

接着，我们用这些 poly_codes 和 bert 的输出做 attention 得到 context_vecs：

def dot_attention(q, k, v, v_mask=None, dropout=None):
  attention_weights = torch.matmul(q, k.transpose(-1, -2))
  if v_mask is not None:
    attention_weights *= v_mask.unsqueeze(1)
  attention_weights = F.softmax(attention_weights, -1)
  if dropout is not None:
    attention_weights = dropout(attention_weights)
  output = torch.matmul(attention_weights, v)
  return output

state_vecs = self.bert(context_input_ids, context_input_masks, context_segment_ids)[0]  # [bs, length, dim]
context_vecs = dot_attention(poly_codes, state_vecs, state_vecs, context_input_masks, self.dropout) #[bs, m, dim]

得到 response_vec 的方式类似，不再赘述。最后，只需根据 response_vec 给 context_vecs 做一次 attention 得到 final_context_vec 即可：

if labels is not None:
      responses_vec = responses_vec.view(1, batch_size, -1).expand(batch_size, batch_size, self.vec_dim)

final_context_vec = dot_attention(responses_vec, context_vecs, context_vecs, None, self.dropout)

在 loss function 方面，虽然我们可以在准备数据的时候就为每个样本做 N 次负采样，但是在 Bi-encoder 或 Poly-encoder 这种产生 response_vec 和 query 完全独立的场景下，可以将同一个 batch 内的其他 response 作为负样本来避免重复计算，有效提升训练效率。

具体实现时，我们计算 context_vec_i 和 response_vec_j 的点乘，从而产生一个 [bs, bs] 的余弦相似度矩阵，这个相似度矩阵就是 context_vec_i 和 batch 内的每一个 response_vec 的相似度。

由于我们的目标是最大化 context_vec_i 和对应的正样本，即 response_vec_i 的相似度，所以我们可以做一个 [bs,bs] 的单位矩阵作为 label，最后应用交叉熵产生训练用的 loss。

我的代码中在 dot_product 后面还乘了系数 5，这就是一个用于缓和 softmax 取值的参数，其具体取值通常需要实验来确定，这里的 5 只是我的经验值。

# 因为要算余弦相似度，所以给向量都归一化一下，之后直接点乘即可
context_vec = F.normalize(context_vec, 2, -1)
responses_vec = F.normalize(responses_vec, 2, -1)

responses_vec = responses_vec.squeeze(1)
dot_product = torch.matmul(context_vec, responses_vec.t())  # [bs, bs]
mask = torch.eye(context_input_ids.size(0)).to(context_input_ids.device)
loss = F.log_softmax(dot_product * 5, dim=-1) * mask
loss = (-loss.sum(dim=1)).mean()

5.2 实验结果

我使用的实验数据是论文中所用的 Ubuntu V2，实验设备是我笔记本上的一个 1066 显卡。当然为了实验跑得更快，我没有使用论文中所用的 bert-base，而是一个预训练过的仅 4 层的 bert-small。

另外，此实验中所用的 batchsize、文本长度、历史对话信息等都限制的比较小（不然实验实在是跑得太慢了），因此实验结果整体会较原论文中偏低。

最终的实验设置和结果如下：

Dataset: Ubuntu V2
Device: GTX 1060 6G x1
Pretrained model: BERT-small-uncased (https://storage.googleapis.com/bert_models/2020_02_20/all_bert_models.zip)
Batch size: 32
max_contexts_length: 128
max_context_cnt: 4
max_response_length：64
lr: 5e-5
Epochs: 3

Results:

▲ 复现实验结果汇总

从上表中明显可以看出，Poly-encoder 的效果要远优于 Bi-encoder 的，当使用 16 个 codes 时，poly 较 bi 的提升可得到 2.24 个点，而使用 64、360 个 codes 时提升分别可达 3.12 和 3.52 个点。而且模型的训练速度几乎没有受到影响，同时对显存的负担也非常小。