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

最新推荐文章于 2024-05-15 16:47:37 发布

rogeroyer

最新推荐文章于 2024-05-15 16:47:37 发布

阅读量985

点赞数 2

分类专栏：论文笔记 Deep Learning 文章标签：语言模型百度人工智能搜索引擎

本文链接：https://blog.csdn.net/roger_royer/article/details/126111799

版权

Deep Learning 同时被 2 个专栏收录

14 篇文章 1 订阅

订阅专栏

论文笔记

4 篇文章 0 订阅

订阅专栏

预训练语言模型在百度网页搜索中的应用

今天刷到了百度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_{cand_{i}}$ ，Query通过 Content Encoder 之后得到 $Out_{x} 1, Out_{x} 2, ..., Out_{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_{c t x t}^{i}=\sum_{j} w_{j}^{c_{i}} h_{j} \quad \text { where } \quad\left(w_{1}^{c_{i}}, ., w_{N_{x}}^{c_{i}}\right)=\operatorname{softmax}\left(c_{i} \cdot h_{1}, . ., c_{i} \cdot h_{N_{x}}\right)$
为了得到最终的context的representation，使用Cand emb 作为Attention中的Q，来整合上述 $m$ 个全局特征向量，计算公式如下：
$y_{c t x t}=\sum_{i} w_{i} y_{c t x t}^{i} \quad \text { Where } \quad\left(w_{1}, . ., w_{m}\right)=\operatorname{softmax}\left(y_{c a n d_{i}} \cdot y_{c t x t}^{1}, . ., y_{c a n d_{i}} \cdot y_{c t x t}^{m}\right)$
最后，该Query和Doc的最终分数是 $y_{c t x t}·y_{cand_{i}}$ ，如Bi-encoder中一样。当 $m < N$ 时，其中 $N$ 表示token的数量，并且poly-encoder只在最顶层进行，这要远比Cross-encoder的自注意力快很多。

接下来，我们回到本文的模型结构，如下图：
model architecture
对比之后可以发现，Cand emb直接使用 [CLS] token的representation $C^{'}$ ，训练和推理阶段的模型结构稍有不同。训练阶段，直接取 $C^{'}$ 和 $P_{1},...,P_{m}$ 内积的最大值；而推理阶段，因为搜索引擎要提前保留Query的特征向量，需要提前计算保存，因此直接对 $P_{1},...,P_{m}$ 求平均池化。

四阶段训练范式

train paradigm

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}$ ，其量化后的索引为：
$I_{i}\left(r_{i}\right)=\left\lfloor\left(r_{i}-s_{i}^{\min }\right) / Q_{i}\right\rfloor$
该索引的范围是 [0, 255]，可以使用8位整型表示。当线上推理时，可以还原其量化前的值为 $\tilde{r_{i}}$ :
$\tilde{r_{i}}=Q I_{i}\left(r_{i}\right) * 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

小结

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