深度学习进阶篇-国内预训练模型[6]：ERNIE-Doc、THU-ERNIE、K-Encoder 融合文本信息和 KG 知识；原理和模型结构详解。

1.ERNIE-Doc: A Retrospective Long-Document Modeling Transformer

1.1. ERNIE-Doc 简介

经典的 Transformer 在处理数据时，会将文本数据按照固定长度进行截断，这个看起来比较”武断”的操作会造成上下文碎片化以及无法建模更长的序列依赖关系。基于此项考虑，ERNIE-Doc 提出了一种文档层级的预训练语言模型方法：ERNIE-Doc 在训练某一个 Segment 时，允许该 segment 能够获得整个 Doc 的信息。

如图 1 所示，假设一篇完整的文档被分割成 3 个 Segment：$S_1,S_2,S_3$，在编码 segment $S_2$时，经典的 Transformer 依赖的只是$S_2$本身：$P(y|S_2)$， Recurrence Transformer (例如 Transformer-XL)依赖的是$S_1,S_2$：$P(y|S_1,S_2)$，这两种方法均没有使 Segment $S_2$获得完整的文档信息$S_1,S_2,S_3$。

图 1 ERNIE-Doc 的建模方式图

但是 ERNIE-Doc 在建模过程中，使得每个 Segment 均能获得完整的文档信息：$P(y|S_1,S_2,S_3)$，其中建模长序列的关键点如下：

• Retrospective feed mechanism: 将文本两次传入模型获得文本序列的 representation，第一次将获得完整的文本序列表示，然后该序列的 representation 继续参与第二次的编码过程，这样该文本序列在第二次编码过程中，每个 token 位置便能获得序列完整的双向信息。

• Enhanced recurrence mechanism: 使用了一种增强的 Segment 循环机制进行建模。

• Segment-reordering objective: 对一篇文档中的各个 segment 随机打乱，获得多个乱序文档，然后要求模型预测这些文档中哪个是正常语序的文档。

1.2. 经典/Recurrence Transformer 的计算

在正式介绍正式 ERNIE-DOC 之前，我们先来回顾一下经典和 Recurrence Transformer 模型的计算。假设当前存在一个长文档$D$​被划分为这样几个 Segment：$D=S_1,S_2,...,S_T$​​，其中每个 Segment 包含$L$​个 token：$S_{\tau }=x_{\tau ,1},x_{\tau ,2},...,x_{\tau ,L}$​​。另外约定$h_{\tau }^n\in {\mathbb{R}}^{L\times d}$为 Transformer 第$n$层第$\tau$​​个 Segment 的编码向量。对于第$\tau +1$​个 Segment $S_{\tau +1}$​，Transformer 第$n$层对其相应的编码计算方式为：

如上述讨论，这两种方式均不能使得每个 Segment 获得 Doc 的完整信息。

1.3. Retrospective feed mechanism

ERNIE-Doc 借鉴了人类阅读的行为习惯，在人类阅读时会分为两个阶段：首先会快速略读一下文档内容，然后回过头来仔细阅读。ERNIE-Doc 基于此设计了 Retrospective feed mechanism，该机制同样包含两个阶段：Skimming phase 和 Retrospective phase。

具体来讲，一篇文档会传入模型两次，第一次被称为 Skimming phase，在该阶段将会获得改文档的完整序列表示。第二次被称为 Retrospective phase，在该阶段将会融入 Skimming phase 获得的完整文档表示，开始进一步的编码计算，从而保证在第二次计算每个 Segment 编码时能够获得完整的文档信息。Retrospective phase 的计算方式如下：

其中以上公式各个参数解释如下：

• $T$​：一篇文档的 Segment 数量； $N$：Transformer 模型层的数量；$L$​：每个 Segment 中的最大 token 数量。

• $\hat{H}\in {\mathbb{R}}^{(L\times T\times N)\times d}$​​​​​：一篇文档在所有层中输出的编码向量。

• ${\hat{H}}_{1:T}^i\in {\mathbb{R}}^{(L\times T)\times d}$ : 一篇文档在第$i$​层产生的编码向量。

• ${\hat{h}}_{\tau }^i$：第$\tau$个 Segment 在第$i$层产生的编码向量。

从以上公式可以看到，在 retrospective 阶段，当计算每个 Segment 时，会引入完整文档的表示$\hat{H}$​，这样就保证了编码时，每个 token 能够获得完整文档的信息。

1.4. Enhanced Recurrence Mechanism

ERNIE-Doc 通过使用 Retrospective feed mechanism 和 Enhanced Recurrence Mechanism 两种方式，增大了计算每个 segment 时的有效上下文长度。但是第 3 节引入的公式计算复杂度是很高，因此 Enhanced Recurrence Mechanism 期望前一个 Segment 便能获得完整的文档信息，然后直接融入前一个 Segment 便能使得当前 Segment 计算融入完整的文档信息。

如图 2 所示，ERNIE-Doc 通过将前一个 Segment 的同层编码表示，引入了当前 Segment 的计算中，这个做法同时也有利于上层信息反补下层的编码表示，具体公式为：

图 2 ERNIE 的 Segment 连接方式

1.5. Segment-Reordering Objective

在预训练阶段，ERNIE-Doc 使用了两个预训练任务：MLM 和 Segment-Reordering Objective。我们先来讨论 Segment-Reordering Objective，其旨在帮助模型显式地建模 Segment 之间的关系，其会将一篇长文档进行划分为若干部分，然后将这些部分进行随机打乱，最后让模型进行预测原始的语序，这是一个$K$​分类问题：$K={\sum }_{i=1}^{m}i!$​，其中$m$​​是最大的划分数量。

如图 3 所示，假设存在一篇文档$D$被划分为 3 部分：$D=C_1,C_2,C_3$，ERNIE-Doc 通过打乱这些部分得到$\hat{D}=C_2,C_3,C_1$，然后在最后一个 Segment $S_{\tau }$的位置进行预测原始的文档顺序$C_1,C_2,C_3$。

图 3 预训练任务 Segment-Reordering Objective

另外，在获得$\hat{D}=C_2,C_3,C_1$后，ERNIE-Doc 会对$\hat{D}$进行划分 Segment：$\hat{D}=S_1,S_2,...,S_T$​​​​，并且会对这些 Segment 中的某些 Token 进行 Mask，从而构造 MLM 任务，要求模型根据破坏的 Segment ${\hat{S}}_{\tau }$恢复成原始的$S_{\tau }$​。结合 MLM 和 Segment-Reordering Objective 总的预训练目标为：

其中，$1_{\tau =T}$表示 Segment-Reordering Objective 仅仅在最后一个 Semgnet $S_T$位置被执行，以优化模型。

• 相关资料

1. ERNIE-Doc: A Retrospective Long-Document Modeling Transformer
   

2. ERNIE-Doc Github
   

2.ERNIE：Enhanced Language Representation with Informative Entities

2.1. THU-ERNIE 简介

当前的预训练模型（比如 BERT、GPT 等）往往在大规模的语料上进行预训练，学习丰富的语言知识，然后在下游的特定任务上进行微调。但这些模型基本都没有使用知识图谱（KG）这种结构化的知识，而 KG 本身能提供大量准确的知识信息，通过向预训练语言模型中引入这些外部知识可以帮助模型理解语言知识。基于这样的考虑，作者提出了一种融合知识图谱的语言模型 ERNIE，由于该模型是由清华大学提供的，为区别百度的 ERNIE，故本文后续将此模型标记为THU-ERNIE。

这个想法很好，但将知识图谱的知识引入到语言模型存在两个挑战：

• Structured Knowledge Encoding：如何为预训练模型提取和编码知识图谱的信息？

• Heterogeneous Information Fusion：语言模型和知识图谱对单词的表示（representation）是完全不同的两个向量空间，这种情况下如何将两者进行融合？

对于第一个问题，THU-ERNIE 使用TAGME提取文本中的实体，并将这些实体链指到 KG 中的对应实体对象，然后找出这些实体对象对应的 embedding，这些 embedding 是由一些知识表示方法，例如TransE训练得到的。

对于第二个问题，THU-ERNIE 在 BERT 模型的基础上进行改进，除了 MLM、NSP 任务外，重新添加了一个和 KG 相关的预训练目标：Mask 掉 token 和 entity (实体) 的对齐关系，并要求模型从图谱的实体中选择合适的 entity 完成这个对齐。

2.2. THU-ERNIE 的模型结构

图 1 THU-ERNIE 的模型架构

THU-ERNIE 在预训练阶段就开始了与 KG 的融合，如图 1a 所示，THU-ERNIE 是由两种类型的 Encoder 堆叠而成：T-Encoder 和 K-Encoder。其中 T-Encoder 在下边堆叠了$N$层，K-Encoder 在上边堆叠了$M$层，所以整个模型共有$N+M$层，T-Encoder 的输出和相应的 KG 实体知识作为 K-Encoder 的输入。

从功能上来讲，T-Encoder 负责从输入序列中捕获词法和句法信息；K-Encoder 负责将 KG 知识和从 T-Encoder 中提取的文本信息进行融合，其中 KG 知识在这里主要是实体，这些实体是通过 TransE 模型训练出来的。

THU-ERNIE 中的 T-Encoder 的结构和 BERT 结构是一致的，K-Encoder 则做了一些改变，K-Encoder 对 T-Encoder 的输出序列和实体输入序列分别进行 Multi-Head Self-Attention 操作，之后将两者通过 Fusion 层进行融合。

2.3. K-Encoder 融合文本信息和 KG 知识

本节将详细探讨 K-Encoder 的内部结构以及 K-Encoder 是如何融合预训练文本信息和 KG 知识的。图 1b 展示了 K-Encoder 的内部细节信息。

我们可以看到，其对文本序列 (token Input) 和 KG 知识(Entity Input)分别进行 Multi-Head Self-Attention(MH-ATT)操作，假设在第$i$层中，token Input 对应的 embedding 是$w_1^{(i-1)},w_2^{(i-1)},...,w_n^{(i-1)}$，Entity Input 对应的 embedding 是$e_1^{(i-1)},e_2^{(i-1)},...,e_n^{(i-1)}$，则 Multi-Head Self-Attention 操作的公式可以表示为：

然后 Entity 序列的输出将被对齐到 token 序列的第一个 token 上，例如实体”bob dylan”将被对齐到第一个单词”bob”上。接下里将这些 MH-ATT 的输入到 Fusion 层，在这里将进行文本信息和 KG 知识的信息融合。因为有些 token 没有对应的 entity，有些 token 有对应的 entity，所以这里需要分两种情况讨论。

对于那些有对应 entity 的 token，信息融合的过程是这样的：

对于那些没有对应 entity 的 token，信息融合的过程是这样的：

其中这里的$\sigma (\cdot )$是个非线性的激活函数，通常可以使用 GELU 函数。最后一层的输出将被视作融合文本信息和 KG 知识的最终向量。

2.4. THU-ERNIE 的预训练任务

在预训练阶段，THU-ERNIE 的预训练任务包含 3 个任务：MLM、NSP 和 dEA。dEA 将随机地 Mask 掉一些 token-entity 对，然后要求模型在这些对齐的 token 上去预测相应的实体分布，其有助于将实体注入到 THU-ERNIE 模型的语言表示中。

由于 KG 中实体的数量往往过于庞大，对于要进行这个任务的 token 来讲，THU-ERNIE 将会给定小范围的实体，让该 token 在这个范围内去计算要输出的实体分布，而不是全部的 KG 实体。

给定 token 序列$w_1,w_2,...,w_n$和对应的实体序$e_1,e_2,...,e_m$，对于要对齐的 token $w_i$来讲，相应的对齐公式为：

类似与 BERT 对 token 的 Mask 策略，THU-ERNIE 在 Mask token-entity 对齐的时候也采用的一定的策略，如下：

1. 以 5%的概率去随机地替换实体，让模型去预测正确的 entity。

2. 以 15%的概率直接 Mask 掉 token-entity，让模型去预测相应的 entity。

3. 以 80%的概率保持 token-entity 的对齐不变，以让模型学到 KG 知识，提升语言理解能力。

最终，THU-ERNIE 的总的预训练损失是由 MLM、NSP 和 dEA 三者的加和。

• 参考资料

1. ERNIE：Enhanced Language Representation with Informative Entities
   

2. ERNIE Githut