腾讯专为实体链接而做的新预训练模型

我家有两只小老虎，它们要祝大家虎年大吉呀，哈哈哈，来呀、上图！

我家小老虎

在放假前刷到Arxiv的一篇和知识图谱相关的论文，是研究实体链接预测的，之前简单扫了一眼，作者都是来自腾讯，想必是经过真实业务场景数据检验的，不过当时没有时间去看，存下来作为 TODO 了，这两天抽点时间大概看了看，简单分享给大家，大佬们不吝赐教~

一、核心纪要

实体链接预测是知识图谱领域很重要的组成部分，而在实体链接预测中发挥很大作用的便是知识表示学习(知识图谱Embedding)，之前有很多的文献调研结果表明知识图谱Embedding可以大致分为两个领域，即基于翻译的距离模型以及语义匹配模型。

基于翻译的距离模型通过设计不同节点之间的距离评估方法，利用基于距离的得分函数来表达节点之间的不同关系。这类方法代表性的模型是以 TransE 为开始的一系列翻译模型及其后续变种，比如 TransE、TransH、RotatE、HAKE等，之前写过一篇文章专门讲述这些翻译模型，感兴趣的同学可以查看。虽然基于翻译距离表示的实体关系距离能够表示的非常多样化，但很难预测尚未出现的实体信息。

第二类方法是语义匹配的方法，这类方法不受冷启动的影响，对于未见的实体表示可以通过文本的上下文获取。这类方法的一些知名代表性模型比如KG-BERT、MLMLM、StAR等。这类方法也有其对应的弊端，训练前阶段只学习了上下文知识，而忽略了关系信息。此外，模型结构通常比较复杂，很难构建高比例的负采样样本，导致训练过程中对负样本信息学习不足。

二、不足改进

针对上文中翻译距离模型中未出现的实体预测能力差以及语义匹配模型训练不足的问题(关系信息、负样本构建)，提出了一种新的Knowledge Graph BERT预训练框架(LP-BERT)，其本质是语义匹配。针对上述的问题，主要有两个部分：

其一采用多任务学习预训练策略，在预训练过程中，不仅仅采用MLM学习上下文知识，而且引入实体语义预测和关系语义预测学习知识图谱中三元组的知识信息，分别为MEM和MRM，它把知识图谱的结构化信息转化为非结构化信息嵌入到预训练过程中。其二受到最近大火的对比学习启发，在一个训练批样本中，加入三元组的负采样方法，在保证训练时间不变的情况下，大大增加了负采样的比例，解决了负采样比例低导致模型训练不足的问题。除此之外，为了进一步提高训练样本的多样性，再次提出了一种基于三元组反向关系的数据扩充方法。

三、模型粗看

LP-BERT 的模型结构主要分为两部分。下图是LP-BERT的整体架构，主要分为多任务预训练阶段(Multi-task pre-training)和知识微调阶段(knowledge finetuning)。多任务预训练任务包含MLM、MEM、MRM三个任务。

整体结构图

1、预训练

下图为多任务预训练的结构图。不同的颜色代表不同的含义，不同的虚线框代表不同的预训练任务。下图中E_h、D_h 分别代表头实体以及头实体对应的文本，R 代表三元组中实体之间的关系，E_t、D_t分别代表尾实体和尾实体对应的文本。X^{[MASK]} 代表预训练中遮掩掉的词，X^{[PAD]} 代表需要补齐的固定长度的向量。MEM_h 代表头实体遮蔽、MEM_t 代表尾实体遮蔽、MRM 代表关系遮蔽，MLM 代表原始BERT中提出的遮蔽语言模型。

预训练框架

Mask Entity Modeling(MEM) ：对于基于语义的实体预测任务，由于每个三元组都包含两个实体：头实体和尾实体，所以针对设计了两个不同的任务:头实体预测和尾实体预测。如上图所示的第一个虚线框为头实体预测，蓝色字体代表头实体的信息，包括遮掩掉的词X^{[MASK]}和真实的label信息 half mile。红色字体代表MLM随机遮蔽掉的词和真实的词。第二个虚线框和第一个一致，只不过是换成了尾实体的遮蔽。

Mask Relation Modeling(MRM) ：对于关系预测任务，样本构建策略类似于MEM任务。在保留三元组中的头尾实体和描述的同时，对关系进行掩码和预测。

Mask Language Modeling(MLM) :为了与MEM和MRM共存，与BERT对序列中所有词进行随机遮蔽预测不同，文中提出的MLM方法只对样本的特定文本范围进行局部随机掩蔽。比如对于头实体预测任务，只会对尾实体(E_t)和尾实体对应的文本(D_t)进行遮蔽和预测，不会影响头实体范围内的信息，其他的尾实体预测以及关系预测都是类似的策略。

Loss Designing：在MEM和MRM任务中构造样本的策略是互斥的，因此同一输入模型训练的三个样本不能同时预测头部实体和尾部实体的预测。为了保证模型的泛化能力，将MEM和MRM任务合并到MIM (Mask Item model)任务中，并定义损失函数如下

2、知识微调

微调阶段主要有两部分，一个是基于对比学习的思想对负采样进行改进，在一批训练样本中构建负样本，能够解决之前方法负样本构建不足，训练不充分的问题。另一个是在训练中采用了数据增强的方法，将原始三元组的关系做了一个反向关系，比如之前的头实体预测样本为 (?,R,E_r) 改写为 (E_r, R_{rev}, ?) 用来增强数据。另外设计了两种距离计算方法来联合计算损失函数，如下：