论文笔记之《Pre-trained Language Model for Web-scale Retrieval in Baidu Search》

今天刷到了百度2021年发表在KDD上的一篇论文《Pre-trained Language Model for Web-scale Retrieval in Baidu Search》，论文中分享了很多干货，包括Query和Doc相关性模型、四阶段训练范式、Embedding压缩和量化、基于ERNIE提升的检索系统workflow、负采样方法等等，并且论文中还给出了算法的伪代码，论文中很多方法与工业界紧密贴合，读完之后，收获多多，所以在此结合自身理解，做一个简要的笔记。

Query-Document相关性模型

该模型架构在Poly-Encoders的基础上做了细微改动以适配搜索场景。Poly-Encoders是Facebook AI 2020年发表在ICLR上的一篇论文，主要在双塔模型的基础上增加了一个Ploy-Attention结构。在介绍文本的模型之前，先简要介绍下Bi-encoder、Cross-encoder、Poly-encoder这三种结构。

Bi-encoder

Bi-encoder其实就是我们常说的双塔模型，左右分别输入Q和D的token_ids，分别经过两个Transformer-base的Encoder，再得到每个token的representation，再使用某种聚合方式（例如：Max-pooling\Mean-pooling）分别得到Q和D的representation（Ctxt Emb和Cand Emb），最后输入两个representation计算得到Score（通常使用内积、余弦相似度等来计算相关性得分）。

Cross-encoder

Bi-encoder有一个弱势在于需要两个Encoder，并且在计算representation时Q和D是独立的，之间没有交互。Cross-encoder就很好地解决了这个问题，将Q和D同时输入一个Encoder，再聚合得到一个representation（Cand Emb），后面再接入全连接神经网络计算得到Score。这种模式就像BERT里面的句子对分类。

Poly-encoder

Poly-encoder综合了Bi-encoder和Cross-encoder的特性，既采用双塔结构，也增加了两个Encoder之间的交互。具体交互由Poly-Encoder模块完成。Doc通过Candidate Encoder编码和 Candidate Aggregator之后得到 Cand emb $y_{can{d}_i}$，Query通过 Content Encoder 之后得到 $[Ou{t}_{x}1,Ou{t}_{x}2,...,Ou{t}_x{N}_x]=[h_1,h_2,...,h_{N_x}]$ 。另外，设置 $m$ 个可学习的向量（context codes）$[c_1,c_2,...,c_m]$作为Attention中的Q，用于计算得到 $m$ 个全局特征向量 $[y_{ctxt}^1,...,y_{ctxt}^m]$，其具体计算方式如下：

$$y_{ctxt}^i=\sum _j{w}_j^{c_i}{h}_j\text{ where }\left(w_1^{c_i},.,w_{N_x}^{c_i}\right)=\mathrm{softmax}\left(c_i\cdot h_1,..,c_i\cdot h_{N_x}\right)$$
为了得到最终的context的representation，使用Cand emb 作为Attention中的Q，来整合上述 $m$ 个全局特征向量，计算公式如下：
$$y_{ctxt}=\sum _i{w}_i{y}_{ctxt}^i\text{ Where }\left(w_1,..,w_m\right)=\mathrm{softmax}\left(y_{can{d}_i}\cdot y_{ctxt}^1,..,y_{can{d}_i}\cdot y_{ctxt}^m\right)$$
最后，该Query和Doc的最终分数是 $y_{ctxt}\cdot y_{can{d}_i}$，如Bi-encoder中一样。当 $m<N$时，其中 $N$ 表示token的数量，并且poly-encoder只在最顶层进行，这要远比Cross-encoder的自注意力快很多。

接下来，我们回到本文的模型结构，如下图：

对比之后可以发现，Cand emb直接使用 [CLS] token的representation $C^{{}^{\prime }}$，训练和推理阶段的模型结构稍有不同。训练阶段，直接取 $C^{{}^{\prime }}$和 $P_1,...,P_m$ 内积的最大值；而推理阶段，因为搜索引擎要提前保留Query的特征向量，需要提前计算保存，因此直接对 $P_1,...,P_m$ 求平均池化。

四阶段训练范式

pretraining-阶段一

随机初始化权重，使用万亿规模的数据（来自中文维基百科、百度新闻、百度百科、百度贴吧）来预训练ERNIE，使用MLM和NSP任务。

post-pretraining-阶段二

采用阶段一的权重，使用万亿规模的搜索日志，采用实体、短语级别的MLM和NSP任务，Doc是否点击作为NSP的label。

intermediate fine-tuning-阶段三

采用阶段二的权重，使用万亿规模的搜索日志（含点击和未点击的Doc），结合本文的模型结构进行微调。

target fine-tuning-阶段四

采用阶段三的权重，使用小规模的更加准确且无偏的人工标注样本（每个Q-D对的得分介于0-4之间），结合本文的模型结构进行微调。

这种多阶段预训练和微调模式确实能在很多场景发挥作用，曾在一个NER竞赛中采用了类似的方法，能带来提升。

Embedding压缩和量化

Embedding压缩

原理就是采用一个全连接神经网络对representation进行降维。

Embedding量化

每个embedding中的值使用float32保存，可以通过转换将其变为uint8，虽然对精度可能有些损害，但大大节省了存储空间。实现细节如下，首先根据大规模验证数据集上的输出embedding，我们计算输出embedding的第 $i$ 维的数据范围为 $(s_i^{min},s_i^{max})$，然后将该数据范围划分为 $L=255$ 份，每一份都是等间隔 $Q_i=(s_i^{max}-s_i^{min})/L$。对于输出embedding上第 $i$ 维的值 $r_i$，其量化后的索引为：
$$Q{I}_i\left(r_i\right)=\lfloor \left(r_i-s_i^{\min }\right)/Q_i\rfloor$$
该索引的范围是 [0, 255]，可以使用8位整型表示。当线上推理时，可以还原其量化前的值为 $\widetilde{r_i}$:
$$\widetilde{r_i}=Q{I}_i\left(r_i\right)\ast Q_i+Q_i/2+s_i^{\min }$$

基于ERNIE提升的检索系统workflow

上图分为两部分，先说左边部分；左边是离线数据库和索引的构建，就是事先计算好每个Doc title的representation，然后存入embedding数据库和建立ANN索引，便于后续使用；右图在早期只有Term匹配的workflow上增加了本文提出的ERNIE提升的workflow，主要在两个地方进行扩充，一是通过Embedding ANN召回了更多相关的Doc，二是在后检索过滤时引入了ERNIE相关性得分这个特征。

模型训练和预测的伪代码

def model_encode_query(tokens): 
	all_embeds = ERNIE_encoder.get_all_outputs(tokens) 
	poly_embeds = poly_attention(all_embeds, context_codes) 
	return [fc_compression(poly_embeds[i]) for i in range(m)]

def model_encode_doc(tokens): 
	cls_embed = ERNIE_encoder.get_cls_output(tokens) 
	return fc_compression(cls_embed)

def train_interaction(q, pos, neg): 
	all_logits1, all_logits2 = [], [] 
	for i in range(m): 
	# in-batch random negative sampling via matrix multiplication 
	pos_logits_with_rand_neg = matmul(q[i], pos.T) 
	neg_logits_with_rand_neg = matmul(q[i], neg.T) 
	all_logits1.append(pos_logits_with_rand_neg) 
	all_logits2.append(neg_logits_with_rand_neg) 
	max_logits1 = reduce_max(all_logits1) 
	max_logits2 = reduce_max(all_logits2) 
	final_logits = concat(max_logits1, max_logits2) 
	loss = softmax_with_cross_entropy(final_logits) 
	return loss

def predict_interaction(q, d): 
	avg_q = reduce_mean(q) 
	return reduce_sum(avg_q * d)

def train(query_tokens, pos_doc_tokens, neg_doc_tokens): 
	query_embed = model.encode_query(query_tokens) 
	pos_doc_embed = model_encode_doc(pos_doc_tokens) 
	neg_doc_embed = model_encode_doc(neg_doc_tokens) 
	loss = train_interaction(query_embed, pos_doc_embed, neg_doc_embed)
	apply_optimization(loss)

def predict(query_tokens, doc_tokens): 
	query_embed = model_encode_query(query_tokens) 
	doc_embed = model_encode_doc(doc_tokens) 
	score = predict_interaction(query_embed, doc_embed) 
	return score
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

小结

本文对论文中提到的Query和Doc相关性模型、四阶段训练范式、Embedding压缩和量化、基于ERNIE提升的检索系统workflow展开了介绍，负采样方法以及损失函数和对比学习方法中常用的比较类似，本文就没展开介绍，感兴趣的可以参考论文《SimCSE: Simple Contrastive Learning of Sentence Embeddings》。总之，本文立足于百度搜索引擎，很具有实践价值。