知识图谱增强的预训练模型

ERNIE算得上是该领域的较早工作，其知识融合方法可通过下图进行表达。

从图中可以看出，ERNIE在传统的transformer encoder（也就是图中的T-encoder，T可以理解为token或者text）的基础上，加入了一种可以进行知识融合的encoder模块（对应图中的K-encoder）。一个K-encoder的输入除了过去常见的token input之外，还有知识的输入。 ERNIE中知识主要指实体，即从待理解的输入语句中提取的一些特定的概念，这些实体可以通过实体识别、匹配等方法找到，当然在学术实验中多是默认为一个已有的值，后续的NLU上的工作基本继承了这一做法。当然，也存在一些既使用实体信息又利用了实体之间的关系信息的工作，后面也会进行介绍。看完输入之后，K-Encoder的输出同样也是包括两部分，即token输出和entity输出。这种设计可以实现将M个K-encoder进行拼接，从而支持进行更深度的语义融合。那么，文本信息和实体信息是怎么进行融合的呢？

之前说过，知识融合的头一个问题是特征空间的问题。所以ERNIE设计了一种projection的方式进行特征融合。简单而言就是，对于同一概念或实体通过文本间self-attention与实体self-attention得到的两种embedding w_j与w_k，通过以下的线性变换加非线性激活函数的方式，先融合到一个新的空间，之后再映射回原有空间。

而对于输入文本中没有对应的实体的那些token，则不需要进行融合，仅仅将上述流程自我变换一遍即可。

在得知了融合的方式之后，另一个问题是：上述融合中的参数，以及对entity进行self-attention部分的参数，如何进行训练？在过往BERT模型提出的MLM与NSP的基础上，此处添加了一个实体对齐任务，类似于MLM，就不赘述了。

这篇工作之后，国内陆续涌现出K-BERT、CoLake等论文，但这些论文同embedding级别的融合工作稍有不同，因此放在后面介绍。下面首先介绍一篇AllenNLP分享的论文。

这篇论文同样是用embedding融合的方式去做的，只不过复杂了一点，废话少说，上图：

图画的还是很乱的，单凭颜色，我们可以粗略把他们分成三块：最右边的蓝色部分是实体，下部红色是融合之前的模型，紫色是融合之后的。文中一共标记了七个步骤，其实是有些繁琐的，此处把这七个步骤简单罗列，感兴趣者可以去看论文。

1. Projection to entity dimension (200)；
2. Mention-span pooling；
3. Mention-span transform；
4. Entity embedding pooling（attention-based）；
5. Entity link；
6. Fusion（recontextualization）；
7. Inv-Projection；

可以看出，以上很多工作都是为融合做准备，且都是ERNIE没有考虑到的问题。如将token归并成span，这是因为有些实体其实包含多个token，换而言之，token和entity不是一一对应的。这一点在中文里尤甚，而ERNIE是没有考虑到的。再比如说，查询知识图谱库时，常常会出现输入一个string出现多个entity的情况，这种不算太学术但实际使用时必须考虑的问题，也被这篇工作考虑到了。除此之外，第一步和第七步的两个逆变换，其实是对ERNIE思路的一种继承。当然，最重要的还是知识融合的过程。这篇论文借用自注意力机制进行融合，具体的融合方式可以用下面的公式进行表达：

我们知道，此处的H是projection之后的token输入，也就是纯粹的没有知识融入的文本特征，而S^’e，是从文本中提取出的span在融合了相关实体信息之后的表示，也就是既包含该实体的上下文信息又包含该实体从知识图谱中获得的额外信息的表示。KnowBERT用H作为Q，用S作为K和V，进行注意力机制的聚合。抛去公式，直观地理解这一过程就是：长度为N的文本中的第i个token与所挖掘M个实体中的第j个实体，通过表示向量的内积计算二者的相似度，而后进行归一化形成注意力矩阵。所以说，同传统的self-attention的方阵不同，此处的注意力矩阵是N*M的，且N一般比M大。比如在上图中，N=4，而M=3。之后，基于这样的一个权重矩阵，每个token的表示向量都通过该token对每一个entity的加权和获得。

可以看出，这种方法不同于ERNIE那种仅是为了让预训练模型理解某几个特定的生词，而是将实体的信息更加显式地插入到了每个token里。当然，这样的结果更多是融合之后的信息，需要通过残差连接维持原有的文本输入。这篇工作的另一个问题是，中间涉及到较多的可学习模块，因此较难训练。一般该类问题都可以通过课程训练，子模块单独训练等方式缓解冷启动问题，此处就不赘述了。 KG-BART 之后再介绍一篇发表在AAAI2021上的论文，KG-BART。同之前的NLU任务不同，KG-BART主要是用在了文本生成上的常识文本生成（commonsense generation）任务。该任务是一种特殊的硬限制文本生成（hard-constraint NLG）,特殊之处就是“常识”。而为了让生成模型具有常识，自然就需要结合知识图谱或相关文档。有关commenGen的相关信息后续还会介绍。这篇论文的整体框图如下图所示：

从图中可以看出，同BART一样，KG-BART也是典型的encoder-decoder结构。该结构的输入是一个BPE编码的token输入，输出为我们期待的句子。可以看出，知识子图在encoder和decoder上都进行了增强。当然，此处的两种知识子图的意义不同，encoder处的图侧重于补充常识信息（R表示reasoning），而decoder部分则主要用以辅助生成，因此是在改变输出token的分布（E表示Expanding）。下面分别对其生成方式进行简单介绍。 Encoder部分： Encoder部分的框图如下图所示：

从图中可以看出，encoder中主要包含三个部分，SCI、MGAT和CSD。其中SCI和CSD类似于know BERT中的span，都是将subword整理成一个concept（或者，entity），以及将这样的一个concept再拆解到token级别的过程。论文中使用的是一维卷积加池化，具体细节此处便不赘述了，让我们集中关注知识的融合过程。 从图中可以看到，知识的融合主要包括实体（v^R）和关系（r^R）两部分，这些embedding的学习均是基于TransE进行的。对于融合过程，该工作基于GAT进行，也就是下面的三个公式：

我们直接看最后一个公式，可以看出这是标准的GAT结构，即图中节点i更新之后的特征等于其邻域内所有节点的特征在经过线性变换后进行有权重的家和。其中，每个节点的初始特征包括两部分，一部分是来自于文本的聚合之后的概念特征，另一部分是知识图谱中该概念对应的embedding在线性变换之后的结果。每个节点的特征即为这两部分的拼接值（也就是第一个公式）。而权重\alpha则是基于目标节点i和邻居节点j的相似度进行计算的（对应第二个公式）。当然，不同于self-attention中使用内积近似相似度，此处使用的是加和形式的相似度，并且额外将两概念之间的关系表示也添加了进去。关于基于知识的图神经网络聚合，还有一些别的聚合方式，相关工作后面会进行介绍。 看完encoder之后，再关注一下decoder。整体而言，decoder的融合形式同know BERT有些类似，下面简单介绍。

上图是decoder部分的框图，同encoder不同的是，该图对知识图谱子图部分的处理更多一些。这是因为该工作认为进行文本生成时需要更多地考虑同当前概念相关的一些token，以更好地进行文本生成。因此，此时的知识图谱子图（右下角）不仅仅包含了encoder输入里的概念（ski，skier，mountatin），还包括了许多新的概念，从而形成几个概念树（每个树的根节点为输入中的一个概念）。而decoder部分先不考虑概念之间的relation，而是首先基于这样的树结构将叶子节点的信息传播到根节点，之后再进行根节点之间的传递和知识与文本信息的融合。对于树上的GNN，论文直接使用了GAT进行特征传播，如下面的公式所示：

对于概念之间的信息传播，同样地一个GAT就可以解决。在此也不赘述。

下面来看一下概念上的表示如何同文本上的概念表示进行融合。正如前面所述，此处使用的是know BERT中的方法，该方法也是transformer decoder中对cross-attention的一个直接的化用。即文本输入作为Q，外部输入作为K和V，如下面公式所展示的：

当然，AT^KG是知识图谱信息与文本信息融合的方法，而AT^TX则是vanilla transformer decoder中已有的部分，最终将二者进行拼接即可。

以上三篇工作均为基于embedding进行融合的典型，几乎也将embedding的结合方式总结得差不多了。当然，由于神经网络的可学习性，有时候模型结构对使用效果的影响并不是唯一的。用知识去增强预训练模型，同模型的输入、训练的任务都息息相关。面向特定的下游任务，不同的融合方式甚至也会产生不同的效果。尽管embedding的方式比较直观，且属于典型的深度学习风格，但该类方法仍然存在以下几类问题：

1. 知识的融合效果受knowledge embedding学习限制；
2. 简单的变换未必能将知识空间与文本空间对齐；
3. 该结构需要对预训练模型重新进行预训练，即使通用的KG+PTM预训练是可行的，针对特定领域的预训练或微调同样存在困难。

下面介绍的第二类方法将提供一些额外的思路。

该类工作中，较早的是K-BERT。K-BERT基于输入序列，通过查询语句中的三元组，将输入语句转化为输入树。之后，通过特殊设计BERT的输入信息编码和mask矩阵，实现推理时模型增强的目的。具体过程可通过下面的框图描述。

通过上图可以看出，对于库克正在北京参观这句话，在其输入到预训练模型中时，token embedding和segment embedding并没有发生变化，唯一发生改变的只是position embedding和mask matrix。对于位置编码，图中展现了两种（红色的和黑色的），通过这两种编码就可以将输入的序列一一对应为一个子图。由于硬编码已经在输入的顺序中体现了出来，因此实际BERT输入的位置编码是描述每一条路径的软编码。而mask矩阵的意图更为简单：输入的自然语言文本是彼此可见的，每个三元组内部也是彼此可见的，但文本中的token同其他三元组是不可见的，这样也就避免了前面所说的知识噪声问题。当然，三元组上的额外信息还是可以间接传入文本中的其他token的。

当然，在这个过程中还有其他一些可做手脚的地方，比如segment embedding，完全可以用它来区分文本输入和知识输入等等。后续也有一些工作顺手在这些地方做了修改。

下面介绍一篇同样是在输入上下文章且更惊艳一点的论文，CoLAKE。

如果说K-BERT是在想，该如何把输入文本扩展成一棵树的话，那么CoLAKE的核心思路就是，能不能把输入文本通过知识图谱扩展成一张图？这当然是可行的。比起K-BERT将输入文本理解为按照位置编码组织起来的一条线，CoLAKE则基于self-attention与GNN相似性上的结论，将输入文本看作是全连接图，并在此基础上将该图通过知识图谱进行扩展。也就形成了下图所示的框图。

构建起这样的图之后，同样地，该图也是需要送入BERT等模型中去的，因而CoLAKE面临着和K-BERT同样的设计问题。下图可以展现出其设计细节：

可以看出，该工作和K-BERT没有太大的区别，主要的区别有三处：1）segment embedding中细化了很多细节；2）该图的mask矩阵是基于图的拓扑关系构建的，相当于是邻接矩阵，而不是根据避免知识噪声的原则；3）该论文更为重视预训练过程，通过之前所说的扩展MLM的方法完成了对模型更好的训练。 CoLAKE的另外一个特色在于，通过这种方式得到的知识图谱中实体与关系的embedding，其实比transE等传统方法的效果还要好，这一点也在一些知识图谱不全的实验中得到了验证。当然，对比的baseline主要是比较经典的一些方法，但也足够惊艳了。

KEPLER是借用此种思路的较简单的工作。其基本框架可以用下文所述的框图进行表达：

从图中可以看出，KEPLER的思想是，使用统一的一个transformer encoder实现对文本信息或实体描述信息的编码，对两个领域分别设计预训练任务，并使用这类预训练任务的联合任务对encoder进行训练。如上图所示，提取文本信息的Encoder同从头部实体与尾部实体中获取embedding的Encoder是一个encoder，因此，通过知识图谱中的实体描述，Encoder可以学习到实体信息，而通过在文本上的预训练，encoder也可以学习到文本层次上的信息，这样就完成了在训练时将知识信息融入到encoder之中。该模型的训练损失主要包括两部分，一部分是常见的MLM，另一部分则是对于Knowledge Embedding的损失，此处主要是使用负采样的损失函数，即

此处的h和t都是正样本，h’和t’为负样本。D_r是同relation相关的距离函数，论文中简单实用了transE的方式，也就是加和距离：

这篇论文整体思路是比较简洁的，同时期的JAKET则设计得更复杂了一些。

JAKET将知识图谱表示学习与知识图谱增强的NLU两个任务视为彼此可以互相促进的任务，并以此思路作为论文的最大卖点。下图展示了JAKET的整体流程。

粗略来看，整张图可分为两种颜色：蓝色和绿色。蓝色部分主要是知识图谱表示学习的部分，而绿色的部分则是知识增强的文本理解模型。先看蓝色部分，同KEPLER类似（其实也是目前做KGE这块都会用到的方式），transformer encoder（即图中的LM1）会被用在提取实体描述的表示上，而这些描述的表示其实就是对应实体的表示。之后，通过适配于知识图谱的GAT，这些特征将依照知识图谱的拓扑结构得到信息传递，并连接在如实体分类、关系预测等预训练任务上。在整个知识图谱表示学习的过程中，值得注意的两个点是：

1. 同KEPLER一样，JAKET也将LM1视为共享的Encoder，同时为了缓解循环依赖问题，LM2还具有额外的encoder LM2。同时，在基于KGE任务进行训练时，LM1是相对而言frozen的（这也是比较合理的，因为LM1用来提取文本表示，属于只不过是被用在了KGE的部分）。

2.不同于KEPLER的是，知识图谱中实体的表示还是会被用来进行LM2的增强。具体的增强过程在介绍绿色部分时介绍。 在理解清楚KGE的基本过程后，下面简单介绍一下GAT部分。同上文中介绍KG-BART里对知识图谱进行GAT的方式有所不同，此处采用的是另一个风格的GAT，具体形式如下所示：

从图中可以看出，聚合方式同样还是GAT，不过有一个地方发生了变化。过去的GAT是将邻居节点的表示E_u经过矩阵映射之后做聚合，而在此处，为了适配知识图谱中的关系的特征，将聚合的对象改成了f(E_u，R_r)=E_u+R_r。这种做法应当也是受到了TransE的影响（E_u+R_r –> E_v）吧。

下面再来看绿色的部分，也就是常见的文本NLU部分的过程。可以看出，该部分首先也是通过LM1获得输入的token级别的初步特征，之后将该特征表示里实体的表示同知识图谱中获得的实体embedding进行融合。此处的融合方式比较简单，论文中直接进行了加和。之后，融合之后的输入序列表示会进入到LM2中，LM2的输出被用来进行预训练模型相关的任务。同之前的工作一样，任务主要包括两个，一个是通常BERT中的MLM，第二个是专门针对于实体的MLM。另外的一个在实现中的细节是，LM1和LM2不一定是两个分离的模型，比如在论文中，二者是一个transformer encoder（论文中使用的是Roberta）的前8层和后8层。

JAKET整体来看比KEPLER复杂了许多，但在实验上有明显的欠缺。或许是提交了不完全的版本吧……