论文标题：Mask and Reason: Pre-Training Knowledge Graph Transformers for Complex Logical Queries
论文地址: https://arxiv.org/abs/2208.07638
论文会议: KDD 2022
17.(2022.8.16)KDD-kgTransformer：复杂逻辑查询的预训练知识图谱Transformer
摘要
知识图(KG)嵌入是不完备KG推理的主流方法。然而，受其固有的浅层和静态结构的限制，它们很难应对日益增长的复杂逻辑查询的，这些复杂逻辑查询包括逻辑运算符、推定边缘、多个源实体和未知的中间实体。在这项工作中，我们提出了具有掩蔽预训练和微调策略的KGTransformer(KgTransformer)。我们设计了一种KG三元组变换方法，使Transformer能够处理KG，并通过混合专家(MoE)稀疏激活进一步加强了KG的处理能力。然后将复杂的逻辑查询描述为掩蔽预测，并引入二阶段掩蔽预训练策略以提高可转移性和泛化能力。在两个基准测试上的广泛实验表明，kgTransformer在九个域内和域外推理任务上的性能始终优于基于KG嵌入的基准模型和先进的编码器模型。此外，kgTransformer可以通过提供完整的推理路径来解释给定的答案，从而进行可解释性推理。
1.引言
知识图(KG)存储和组织人类关于现实世界的知识，如人工管理的Freebase和Wikidata，以及半自动构建的知识图NELL和Knowledge Vault。在KG的发展过程中，知识表示学习是查询KG的基本问题之一。它的主要挑战在于知识不完备和查询效率低下。众所周知，大规模的KG存在链接缺失的问题，而SPARQL等专门的查询工具无法很好地处理这一问题。
知识图嵌入(KGE)旨在将实体和关系嵌入到低维连续向量中，在过去的十年中蓬勃发展。具体地说，KGE在简单的知识图补全(KGC)问题 ( h , r , ? ) (h,r,?) (h,r,?)中发现了广泛的采用，该问题的特征是单个头部实体 h h h、关系 r r r和缺失的尾实体。然而，现实世界中的查询可能会更加复杂，包括输入边、多个源实体、存在正类一阶(EPFO)逻辑和未知中间节点，即复杂的逻辑查询。
图1展示了复杂逻辑查询“明尼苏达州出生的诺贝尔奖获得者演奏了什么乐器？”的分解查询图。该查询由两个源实体(“明尼苏达”和“诺贝尔奖”，记为蓝色小球 ? \bullet ?)、一阶逻辑运算符合取(∧)、未知中间实体(“在明尼苏达出生并获得诺贝尔奖的人”，记为灰色小球 ? \bullet ?)、未知目标实体(“他们可能演奏的乐器”，记为绿色小球 ? \bullet ?)和潜在的缺失边组成。回答这类查询的一个主要挑战是其不断增长的跳数和逻辑运算符组合的指数复杂性。此外，在推理过程中还应考虑单个查询中围绕多个实体和关系的丰富上下文信息。
图1-EPFO查询推理：基于KGE的推理模型 vs. 预训练的kgTransformer。掩蔽预测训练可以赋予GNN回答EPFO查询的自然能力。
因此，这种查询超出了现有基于KGE的方法的能力范围。首先，大多数KGE的体系结构都很浅，并且使用静态向量，这限制了它们捕获大规模模式的表达能力和能力。其次，恢复一阶缺失链的训练目标不符合复杂逻辑查询的高阶图性质，因此KGE在不训练样本监督数据集上的辅助逻辑函数的情况下无法处理复杂查询。此外，现有的KG基准数据集通常由有限的特定于任务的查询类型组成，这实际上阻止了KGE方法很好地推广到域外查询。
1.1 贡献
在这项工作中，我们建议从体系结构和训练目标两个角度学习用于回答复杂查询的深层知识表示。
在体系结构级别，目标是设计一个为KG指定的深层模型，以便能够处理复杂的逻辑查询。
在训练目标层面，预计该模型将从KG学习通用知识，而不是特定于任务的知识，从而使该模型能够对域外查询具有很强的泛化能力。
为此，我们提出了kgTransformer，这是一个基于Transformer的GNN体系结构，使用自我监督的预训练策略，用于处理复杂的逻辑查询。
1.2 kgTransformer
我们开发了kgTransformer来对KG进行编码。具体地说，为了表示KG中的关系，我们提出了==三元组转换策略，将关系转化为关系节点，从而将KG转换为没有边属性的有向图==。
为了以较低的计算代价进一步扩大模型容量(扩大参数)，我们采用了MoE策略来利用Transormer前馈层的稀疏激活特性。总而言之，这些策略使kgTransformer架构具有高容量和计算效率，使其能够回答关于KG的EPFO查询。
1.3 掩蔽预训练
为了进一步提高kgTransformer的泛化能力，我们引入了掩蔽预训练框架对其进行训练。我们将复杂的逻辑查询回答问题描述为一个掩蔽预测问题。在预训练过程中，我们从KG中随机抽样子图，并掩蔽随机实体进行预测。它包括两个连续的密集初始化阶段和稀疏细化阶段，密集初始化的目标是通过在密集和任意形状的上下文中训练来丰富模型，稀疏细化是在稀疏的小规模元图上训练的，以缓解预训练和下游查询之间的差距。
在两个广泛使用的KG基准测试(即FB15k-237和NELL995)上的广泛实验表明，在九个域内和域外下游推理挑战上，kgTransformer的性能优于最先进的方法，特别是基于KGE的方法。此外，案例研究表明，隐蔽的预培训可以通过提供对未知中间体的预测来赋予kgTransformer的推理可解释性和可解释性。
2.EPFO逻辑查询
我们在KG上引入了存在正一阶(EPFO)逻辑查询，并确定了其挑战。
2.1 EPFO
设
G
=
(
E
,
R
)
\mathcal{G}=(\mathcal{E},\mathcal{R})
G=(E,R)表示KG，其中
e
∈
E
e∈\mathcal{E}
e∈E表示实体，
r
∈
R
r∈\mathcal{R}
r∈R是二元谓词(或关系)
r
r
r：
E
×
E
→
{True,False
}
\mathcal{E} \times \mathcal{E} \rightarrow \text{\{True,False}\}
E×E→{True,False}，其指示关系是否对一对实体成立。给定一阶逻辑存在(?)和合取(∧)操作，合取查询被定义为：
Q
[
A
]
?
?
A
:
?
E
1
,
…
,
E
m
.
e
1
∧
…
∧
e
n
?where?
e
i
=
r
(
c
,
E
)
,
?with?
E
∈
{
A
,
E
1
,
…
,
E
m
}
,
c
∈
E
,
r
∈
R
?or?
e
i
=
r
(
E
,
E
′
)
,
?with?
E
,
E
′
∈
{
A
,
E
1
,
…
,
E
m
}
,
E
≠
E
′
,
r
∈
R
(1)
\begin{aligned} &Q[A] \triangleq ? A: \exists E_1, \ldots, E_m . e_1 \wedge \ldots \wedge e_n \\ &\quad \text { where } e_i=r(c, E), \text { with } E \in\left\{A, E_1, \ldots, E_m\right\}, c \in \mathcal{E}, r \in \mathcal{R} \\ &\quad \text { or } e_i=r\left(E, E^{\prime}\right), \text { with } E, E^{\prime} \in\left\{A, E_1, \ldots, E_m\right\}, E \neq E^{\prime}, r \in \mathcal{R} \end{aligned}\tag{1}
?Q[A]??A:?E1?,…,Em?.e1?∧…∧en??where?ei?=r(c,E),?with?E∈{A,E1?,…,Em?},c∈E,r∈R?or?ei?=r(E,E′),?with?E,E′∈{A,E1?,…,Em?},E=E′,r∈R?(1)
其中
A
A
A指的是查询的(未知)目标实体，
E
1
,
?
?
,
E
m
E_1,\cdots,E_m
E1?,?,Em?指的是指存在量化的绑定变量(即未知的中间实体集)，
c
c
c是指源实体。给定查询
Q
\mathcal{Q}
Q，目标是找到其目标实体集合
A
?
E
\mathcal{A}?\mathcal{E}
A?E满足
a
∈
A
a \in \mathcal{A}
a∈A当且仅当
Q
[
a
]
\mathcal{Q}[a]
Q[a]为真。
对于上述的逻辑查询“What musical instrucments did Minnesota-born Nobel Prize winners play?”，解析得到的EPFO查询蕴含了一个若干约束的二跳(2-hop)知识图谱查询：第一个约束是查询一个出生在Minnesota的实体；第二个约束是查询一个获得了诺贝尔奖的实体，与上一个约束指向同一个实体；第二跳是查询上面得到的实体所Play的乐器。可知，第二跳的查询结果才是问句所想要得到的答案。
除了合取查询(∧)外，EPFO还包括析取查询(∨)。经验法则是将EPFO查询转换为析取范式。换句话说，一个析取查询可以分解成几个合取查询，这样就可以应用一个规则来合成析取预测的合取结果。
2.2 挑战
与传统的基于KGE方法的KGC任务相比，EPFO查询可以是多跳的，众多的EPFO查询组合测试了推理模型的域外泛化能力。所有这些特征共同给推理模型带来了以下挑战：
- 指数复杂性：EPFO查询的复杂性随着跳数的增加呈指数增长，需要高容量和高级模型来处理它们。基于KGE的推理模型依靠嵌入和简单运算符以“从左到右”的自回归方式进行推理。然而，有证据表明，随着嵌入维度增加到1000，这类模型的性能逐渐饱和。此外，在推理过程中，第一个编码的实体不知道后来编码的实体，忽略了有用的双向交互。
- 转移和泛化：经过训练后，理想的推理模型应该会转移到域外查询并进行泛化。但是现有的EPFO推理模型以有监督的方式在几个查询类型(即图2(b)中的1p、2p、3p、2i和3i)内的有限数量的样本上被直接训练，使得原始KG中的许多实体和关系保持不变。因此，推理模型被限制掌握超出现有类型所能表达的各种形式和更大上下文的知识，从而损害泛化能力。
总之，这些挑战使得EPFO查询不同于传统的KGC，后者只涉及单跳和单一类型的查询。在这项工作中，我们探索了如何通过预训练和微调范例有效地处理EPFO查询。
3.KG预训练框架
在本节中，我们将介绍kgTransformer，这是一种基于Transformer的图神经网络(GNN)，用于处理KG上的EPFO查询。为了使kgTransformer具有较强的泛化能力，我们设计了一个掩蔽的预训练和微调框架。整个KG预训练和推理框架如图2所示。
3.1 KgTransformer架构
如上所述，基于KGE的推理模型的相对简单的体系结构限制了它们的表达能力。为了解决这一问题，我们提出了一种基于Transformer的GNN体系结构kgTransformer。用来对KG编码。
为了表示KG中的关系，我们提出了三元组转换策略，将关系转化为关系节点，从而将KG转换为没有边属性的有向图。
为了扩大模型容量(扩大参数)，我们采用了MoE策略，同时保持其计算效率。
3.1.1 用于KG的Transformer
Transformer是一种最初提出用于处理序列的神经体系结构，在图形领域取得了早期的成功。要将Transformer应用于KG，有两个问题需要回答：1)如何在图结构中编码节点邻接关系；2)如何对实体和关系进行建模。
首先，在图形Transformer中编码邻接节点已经有了很好的实践。对于节点和图的分类，通常将图视为具有位置编码的标记序列，而忽略邻接矩阵。然而，对于链路预测，邻接矩阵可能是至关重要的，并且应该掩蔽到自注意力以获得更好的性能。由于==EPFO推理本质上是一个链接预测问题==，所以我们模仿HGT的方法，将邻接矩阵掩蔽到自注意力。
其次，到目前为止，关于如何将实体和关系都纳入Transformer的计算中的研究很少。在这里，我们提出了三元组变换操作，记为函数 L ( ? ) L(\cdot) L(?)，将关系转化为关系节点，从而将KG转换为没有边属性的有向图。
对于 G \mathcal{G} G中的每个有向三元组 ( h , r , t ) (h,r,t) (h,r,t)，其中 r ( h , t ) = True r(h,t)=\text{True} r(h,t)=True，我们在 L ( G ) L(\mathcal{G}) L(G)中创建一个关系节点 r h t r_{ht} rht?，它同时连接到 h h h和 t t t。通过将 G \mathcal{G} G中的每条关系边映射到关系节点，结果图变成 L ( G ) = ( E ′ ， R ′ ) L(\mathcal{G})=(\mathcal{E}'，\mathcal{R}') L(G)=(E′，R′)，其中 E ′ = E ∪ { r h t ∣ r ( h , t ) = True } \mathcal{E}'=\mathcal{E}∪\{r_{ht}|r(h,t)=\text{True}\} E′=E∪{rht?∣r(h,t)=True}和非属性边集 R ′ = { r : r ( h , r h t ) = r ( r h t , t ) = True ∣ r ( h , t ) = True } \mathcal{R}'=\{r:r(h,r_{ht})=r(r_{ht},t)=\text{True}|r(h,t)=\text{True}\} R′={r:r(h,rht?)=r(rht?,t)=True∣r(h,t)=True}。在实践中，三元组变换的计算成本很低，因为EPFO查询的推理图通常非常小且稀疏(~10个实体和关系)。
即： ( h , r , t ) ? h → r , r → t (h,r,t) \Longrightarrow h \rightarrow r , r \rightarrow t (h,r,t)?h→r,r→t
给定一个输入(三元组转换成的有向子图)的嵌入
x
e
(
0
)
∈
R
d
\mathbf{x}^{(0)}_e \in \mathbb{R}^d
xe(0)?∈Rd的
e
∈
E
′
e \in \mathcal{E}'
e∈E′的维为
d
d
d，我们增加了一个特殊的节点类型嵌入
t
I
(
e
∈
E
)
t_{\mathbb{I}(e \in \mathcal{E})}
tI(e∈E)?来区分它是实体节点还是关系节点。在kgTransformer的第
k
k
k层中，
d
h
=
d
/
H
d_h=d/H
dh?=d/H，
H
H
H表示注意力头数，多头注意力的计算公式为：
A
t
t
n
i
(
k
)
=
s
o
f
t
m
a
x
(
Q
T
K
/
d
h
)
V
(2)
\mathrm{Attn}_i^{(k)}=\mathrm{softmax}(Q^TK/\sqrt{d_h})V \tag{2}
Attni(k)?=softmax(QTK/dh??)V(2)
其中 Q = x e ( k ? 1 ) W Q , { K , V } = ∣ ∣ N ∈ N ( e ) x n ( k ? 1 ) { W K , W V } Q=\mathbf{x}^{(k-1)}_e\mathbf{W}_Q,\{K,V\}=||_{N \in \mathcal{N}(e)} \mathbf{x}^{(k-1)}_n \{\mathbf{W_K},\mathbf{W_V}\} Q=xe(k?1)?WQ?,{K,V}=∣∣N∈N(e)?xn(k?1)?{WK?,WV?}(KV表示相邻节点信息)，这里的 W { Q , K , V } ∈ R d × d h , ∣ ∣ \mathbf{W_{\{Q,K,V\}}} \in \mathbb{R}^{d \times d_h},|| W{Q,K,V}?∈Rd×dh?,∣∣表示连接， N ( e ) \mathcal{N}(e) N(e)指的是节点 e e e的邻居集合。(即：从第k-1层的QKV到第k层的Attn)
和Attention Is All You Need一样，在Encoding的时候，通过多头注意力机制来计算得到每个节点的相邻节点与它的关系重要程度，并通过FFN来进行变换得到最终输出。
接下来，我们在由
W
O
∈
R
d
×
d
\mathbf{W_O} \in \mathbb{R}^{d \times d}
WO?∈Rd×d投影的注意力权重输出上应用前馈神经网络FFN(
x
x
x)：
x
e
(
k
)
=
F
F
N
(
∥
i
=
1
H
Attn
?
i
(
k
)
?
W
O
)
,
F
F
N
(
x
)
=
σ
(
x
W
1
+
b
1
)
W
2
+
b
2
(3)
\mathbf{x}_e^{(k)}=\mathrm{FFN}\left(\|_{i=1}^H \operatorname{Attn}_i^{(k)} \cdot \mathbf{W}_O\right), \mathrm{FFN}(\mathbf{x})=\sigma\left(\mathbf{x} \mathbf{W}_1+\mathbf{b}_1\right) \mathbf{W}_2+\mathbf{b}_2 \tag{3}
xe(k)?=FFN(∥i=1H?Attni(k)??WO?),FFN(x)=σ(xW1?+b1?)W2?+b2?(3)
其中
W
1
∈
R
d
×
4
d
\mathbf{W_1} \in \mathbb{R}^{d \times 4d}
W1?∈Rd×4d，
W
2
∈
R
4
d
×
d
\mathbf{W_2} \in \mathbb{R}^{4d \times d}
W2?∈R4d×d，
σ
\sigma
σ是激活函数GeLU。简单来说，FFN是由两个全连接层构成的网络，即Linear(ReLU(Linear(x?)))，其中Linear(x)=xW+b，FNN不改变矩阵大小，FFN对于Transformers捕获大规模模式至关重要，并且已被证明等同于键-值网络。
即==kgTransformer共有L层，从k-1层到k层：第k-1层的输入嵌入通过多头注意力生成第k-1层的QKV，然后QKV经过FFN生成第k层的输入嵌入。==
不同于现有的基于KGE的推理模型具有从左到右的推理顺序，也不同于基于序列编码器的模型只能对非循环查询图进行推理，kgTransformer设计了一种在每一层的计算中聚合来自各个方向的信息的体系结构(Transformer相比于的RNN的优势即为)，使其更加灵活，能够回答任意形状的查询。图2(a)说明了kgTransformer的体系结构。
文中所使用的Transformer大体结构及其内部公式与 Attention Is All You Need是一致的，只是sequence中每个token从原来的单词文字变成了现在的节点。
有一点与Attention Is All You Need不一样的是，对于前馈神经网络(FFN)，作者初始化了多个FFN来进行学习，以模型专家投票的机制，称为Mix-of-Expert(MoE)。
图2(a)：kgTransformer和MoE的结构
3.1.2 混合专家(MoE)模型
MoE在不改变FFN计算量的情况下，扩大了模型容量(扩大参数)(先切割成小的矩阵，计算量呈开方减小，然后再求和)
虽然kgTransformer的架构允许其捕获复杂的推理模式，但其参数的数量随着嵌入维度 d d d增大沿二次曲线飙升，这是Transformers面临的常见挑战。
如前所述，Transformer的FFN被认为等同于键-值网络，即FFN通过全连接的方式实现了与注意力机制类似的效果，其全连接的权重可以从侧面被看做是输入间的重要性关系，因此FFN输出的元素里大部分都是0，是很稀疏的（也因为存在稀疏性，作者认为使用MoE，把一个很大的FFN分解成若干个小的FFN，能够更好的增强模型泛化能力）。
在该网络中，键激活一些值作为响应。给定
x
∈
R
d
x \in \mathbb{R}^d
x∈Rd,FFN的中间激活值
σ
(
x
W
1
+
b
1
)
=
[
x
0
,
x
1
,
…
,
x
i
,
…
,
x
j
,
…
,
x
4
d
]
=
[
0
,
0
,
…
,
x
i
,
…
,
x
j
,
…
,
0
]
?
Most?of?elements?are?
0
(4)
\sigma\left(\mathbf{x W}_1+\mathbf{b}_1\right)=\left[\mathbf{x}_0, \mathbf{x}_1, \ldots, \mathbf{x}_i, \ldots, \mathbf{x}_j, \ldots, \mathbf{x}_{4 d}\right]=\underbrace{\left[0,0, \ldots, \mathbf{x}_i, \ldots, \mathbf{x}_j, \ldots, 0\right]}_{\text {Most of elements are } 0} \tag{4}
σ(xW1?+b1?)=[x0?,x1?,…,xi?,…,xj?,…,x4d?]=Most?of?elements?are?0
[0,0,…,xi?,…,xj?,…,0]??(4)
通常非常稀疏，其中
W
1
∈
R
d
×
4
d
\mathbf{W}_1 \in \mathbb{R}^{d \times 4d}
W1?∈Rd×4d。稀疏性的程度因任务而异。例如，最近的一项研究表明，在NLP的每一次输入中，通常只有不到5%的神经元被激活。在我们对EPFO查询的初步实验中(参见图3)，对于某些输入，只有10%-20%的神经元被激活(最后一层解码器层除外)。
图3：初步实验中kgTransformer沿预训练步骤(无MoE)的FFN激活率。
因此，我们建议通过混合专家(MoE)策略来利用kgTransformer的稀疏性。MoE首先将一个较大的FFN分解为分块专家模型，然后利用光门网络选择专家模型参与计算：
-
例如，具有 W 1 ∈ R d × 16 d \mathbf{W}_1∈\mathbb{R}^{d \times \mathbf{16d}} W1?∈Rd×16d和 W 2 ∈ R 16 d × d \mathbf{W}_2∈R^{\mathbf{16d}×d} W2?∈R16d×d(比公式(4)大4倍)的FFN可以变换有8部分的分块专家模型为
FFN ? ( x ) = { FFN ? Exp ? ( i ) ( x ) = σ ( x W 1 ( i ) + b 1 ) W 2 ( i ) + b 2 ∣ i ∈ [ 1 , 8 ] , i ∈ N } (5) \operatorname{FFN}(\mathbf{x})=\left\{\operatorname{FFN}_{\operatorname{Exp}}^{(i)}(\mathbf{x})=\sigma\left(\mathbf{x} \mathbf{W}_1^{(i)}+\mathbf{b}_1\right) \mathbf{W}_2^{(i)}+\mathbf{b}_2 \mid i \in[1,8], i \in \mathbb{N}\right\} \tag{5} FFN(x)={FFNExp(i)?(x)=σ(xW1(i)?+b1?)W2(i)?+b2?∣i∈[1,8],i∈N}(5)
其中 W 1 ( i ) ∈ R d × 2 d , W 2 ( i ) ∈ R 2 d × d \mathbf{W}_1^{(i)} \in \mathbb{R}^{d \times \mathbf{2d}}, \mathbf{W}_2^{(i)} \in \mathbb{R}^ {\mathbf{2d} \times d} W1(i)?∈Rd×2d,W2(i)?∈R2d×d。 每个 F F N Exp ? ( i ) \mathrm{FFN}_{\operatorname{Exp}}^{(i)} FFNExp(i)? 都被称为专家模型。 -
给定一个门控网络，通常是一个可训练矩阵 W gate? ∈ R d × N \mathbf{W}_{\text {gate }} \in \mathbb{R}^{d \times N} Wgate??∈Rd×N，其中 N N N为块专家模型的数量， 我们可以选择前2位块专家模型进行计算：
x e ( k ) = ∑ ( i , s ) ∈ S s ? F F N Exp ? ( i ) ( x e ( k ? 1 ) ) (6) \mathbf{x}_e^{(k)}=\sum_{(i, s) \in S} s \cdot \mathrm{FFN}_{\operatorname{Exp}}^{(i)}\left( \mathbf{x}_e^{(k-1)}\right) \tag{6} xe(k)?=(i,s)∈S∑?s?FFNExp(i)?(xe(k?1)?)(6)
其中系数s是为了保证求和的“总权重值”为1， S = { ( i , s ) ∣ s ∈ softmax ? ( S=\left\{(i, s) \mid s \in \operatorname{softmax}\left(\right.\right. S={(i,s)∣s∈softmax(Top-2( x e ( k ? 1 ) W gate ) ) \left.\left.\mathbf{x}_e ^{(k-1)} \mathbf{W}_{\text {gate}}\right)\right) xe(k?1)?Wgate?))，其中i是s的索引。
值得注意的是：在训练时，门控机制只会选择权重最大的两个FFN，将其输出进行求和并作为最终输出，以保证训练过程中的时间不过长；但是==在预测时，不使用门控机制，将所有的FFN结果求和并作为最终输出。==
通过使用MoE，我们可以用与公式(4)中FFN大致相同的训练计算量显著扩大实际模型大小。在推理过程中，所有专家都被包括在计算中，FFN输出成为每个专家输出的加权平均。
MoE首先将一个较大的FFN分解为分块专家模型，然后利用光门网络选择几个块中的2个块专家模型参与计算(选择最好的两个 x e ( k ? 1 ) W gate \mathbf{x}_e ^{(k-1)} \mathbf{W}_{\text {gate}} xe(k?1)?Wgate?，把对应的 F F N ( i ) \mathrm{FFN}^{(i)} FFN(i)求和)。
几个块专家模型只有W1和W2矩阵的具体值不同。
3.2 掩蔽预训练和微调
为了进一步提高kgTransformer的泛化能力，我们提出了一个用于KG推理的掩蔽预训练和微调框架，如图2(b)所示。
3.2.1 两阶段的预训练：初始化与优化
大多数现有的EPFO推理模型以有监督的方式在有限数量的采样查询而不是特定的查询类型上训练嵌入。训练中原始KG中实体和关系的覆盖率较低，导致测试中对于未知实体的查询可转换性较差。此外，推理模型通常将在训练后回答域外查询。例如，标准基准测试针对5种类型的查询(1p、2p、3p、2i和3i)训练模型，并要求模型测试其他类型的查询(即域外查询)(ip、pi、2u和up)。
为了克服这一挑战，我们建议为kgTransformer进行掩蔽预训练。其主要思想是从原始KG的训练集中随机抽样任意形状的掩蔽子图来预先训练模型。它的优势在于采样子图的接受范围大，查询形状多样，对原始KG的覆盖率高。为了充分挖掘预训练的通用知识，我们应该尽可能大而密集地对掩蔽查询图进行采样，而在下游推理过程中，查询图通常是小而稀疏的(即规模小，且边或关系很少)。因此，这种策略将导致预训练和下游任务之间的不匹配，导致性能下降。
为了缓解这个问题，我们为kgTransformer引入了两阶段掩蔽预训练策略。第一阶段利用KG的常识对kgTransformer进行掩蔽初始化，第二阶段进一步提高其在推理过程中处理规模小且稀疏的查询的能力。
3.2.1.1 阶段1：密集初始化(完形填空)
这一阶段的目标是通过对密集和大规模的采样子图进行掩蔽预训练，使kgTransformer具备KG中的通用知识。
我们使用两种基于随机游走的策略对原始KG训练集进行采样，可以得到随机不同的子图。一种是带重启的随机游走(RWR)，该策略可以采样任何形状的子图；另一种是基于树的RWR策略，对于该策略，我们将采样图的形状约束为树结构，以满足合取(∧)查询(即专门加入 tree-base是因为EPFO的逻辑连接关系更像树形结构)。对于更密集的上下文，我们收集了采样实体之间的大多数推导关系。
这两种方法都用于生成采样查询子图。给定由
σ
\sigma
σ参数化的kgTransform，由特殊掩码嵌入掩蔽的一组随机采样实体
E
m
a
s
k
\mathcal{E}_{mask}
Emask?将可训练的输入嵌入
X
=
[
x
e
0
(
0
)
,
?
?
,
x
e
T
(
0
)
]
\mathbf{X}=[\mathbf{x}^{(0)}_{\mathbf{e}_0},\cdots,\mathbf{x}^{(0)}_{\mathbf{e}_T}]
X=[xe0?(0)?,?,xeT?(0)?]变为掩码损坏的
X
^
\hat{\mathbf{X}}
X^。密集初始化阶段就寻求优化损失函数：
L
i
n
i
t
=
?
∑
e
m
∈
E
mask?
log
?
p
θ
(
x
e
m
∣
X
^
,
A
)
(7)
\mathcal{L}_{\mathrm{init}}=-\sum_{e_{\mathrm{m}} \in \mathcal{E}_{\text {mask }}} \log p_\theta\left(\mathbf{x}_{e_{\mathrm{m}}} \mid \hat{\mathbf{X}}, \mathbf{A}\right) \tag{7}
Linit?=?em?∈Emask??∑?logpθ?(xem??∣X^,A)(7)
其中
A
\mathbf{A}
A是邻接矩阵。以非自回归的方式对所有被掩蔽的节点同时进行预测。
这里的初始化和BERT的完形填空一样，训练模型预测所有的Mask单词，使得loss最小。
3.2.1.2 阶段2：稀疏优化
此阶段的目标是通过对EPFO查询使用元图采样来进一步增强模型能力，其中EPFO查询通常是小规模且稀疏的。
我们采用最基本的EPFO模式(1p、2p、3p、2i和3i)作为元图形，它涉及的实体和关系不超过四个(模拟规模小的图)，其余所有实体之间的推导关系都不保留(模拟边或关系少的图)。按照实际的查询实践，除了源实体(图2(b)中的“蓝色小球 ? \bullet ?”)之外，其他实体都被掩蔽，但在每种查询类型中只预测目标实体(图2(b)中的“绿色小球 ? \bullet ?”)。
给定每个元图类型，掩码实体集合
E
m
a
s
k
=
E
i
n
t
e
r
∪
E
t
a
r
g
e
t
\mathcal{E}_{mask}=\mathcal{E}_{inter}∪\mathcal{E}_{target}
Emask?=Einter?∪Etarget?，其中
E
i
n
t
e
r
\mathcal{E}_{inter}
Einter?和
E
t
a
r
g
e
t
\mathcal{E}_{target}
Etarget?分别指中间实体和目标实体的集合。稀疏优化阶段仅优化目标实体
e
t
∈
E
t
a
r
g
e
t
e_t∈\mathcal{E}_{target}
et?∈Etarget?上的误差
L
r
e
f
i
n
e
\mathcal{L}_{refine}
Lrefine?，计算形式与
L
i
n
i
t
\mathcal{L}_{init}
Linit?相同：
L
refine?
=
?
∑
e
t
∈
E
target?
log
?
p
θ
(
x
e
t
∣
X
^
,
A
)
(8)
\mathcal{L}_{\text {refine }}=-\sum_{e_t \in \mathcal{E}_{\text {target }}} \log p_\theta\left(\mathbf{x}_{e_{\mathrm{t}}} \mid \hat{\mathbf{X}}, \mathbf{A}\right) \tag{8}
Lrefine??=?et?∈Etarget??∑?logpθ?(xet??∣X^,A)(8)
在kgTransformer预训练过程中，这两个阶段是依次进行的，而不是同时并行进行。
和微调一样，只针对目标实体进行预测训练。
3.2.2 微调
表2：预训练和微调的消融(Hits@3m)
?FB15k-237?
?NELL995?
?kgTransformer?(Stage?1+?Stage?2)?
0.336
0.395
-only?Stage?1?in?pre-training?
0.308
0.307
-only?Stage?2?in?pre-training?
0.307
0.399
-w/o?fine-tuning?
0.301
0.368
-w/o?pre-training?
0.262
0.288
\text{表2：预训练和微调的消融(Hits@3m)}\\ \begin{array}{ccc} \hline & \text { FB15k-237 } & \text { NELL995 } \\ \hline \text { kgTransformer (Stage 1+ Stage 2) } & \mathbf{0 . 3 3 6} & 0.395 \\ \hline \text {-only Stage 1 in pre-training } & 0.308 & 0.307 \\ \text {-only Stage 2 in pre-training } & 0.307 & \mathbf{0 . 3 9 9} \\ \hline \text {-w/o fine-tuning } & 0.301 & 0.368 \\ \text {-w/o pre-training } & 0.262 & 0.288 \\ \hline \end{array}
表2：预训练和微调的消融(Hits@3m)?kgTransformer?(Stage?1+?Stage?2)?-only?Stage?1?in?pre-training?-only?Stage?2?in?pre-training?-w/o?fine-tuning?-w/o?pre-training???FB15k-237?0.3360.3080.3070.3010.262??NELL995?0.3950.3070.3990.3680.288??
我们介绍了kgTransformer针对下游推理任务的微调策略。虽然在预训练中存在针对target进行专门训练的稀疏优化阶段，但我们的初步实验(参见表2)显示kgTransformer的微调仍然是必要的。我们用标记的监督信息来为下游任务调整预训练模型。
具体地说，我们使用掩蔽预测形式的下游EPFO推理数据集来微调预训练的kgTransformer。给定掩码查询图的一组可行答案
E
A
\mathcal{E}_A
EA?，每个查询的目标节点的损失函数被表示为：
L
f
t
=
?
1
∣
E
A
∣
∑
e
a
∈
E
A
log
?
exp
?
(
x
e
a
(
L
)
?
u
e
a
)
exp
?
(
x
e
a
(
L
)
?
u
e
a
)
+
∑
e
?
E
A
exp
?
(
x
e
(
L
)
?
u
e
)
(9)
\mathcal{L}_{\mathrm{ft}}=-\frac{1}{\left|\mathcal{E}_A\right|} \sum_{e_a \in \mathcal{E}_A} \log \frac{\exp \left(\mathbf{x}_{e_a}^{(L) \top} \mathbf{u}_{e_a}\right)}{\exp \left(\mathbf{x}_{e_a}^{(L) \top} \mathbf{u}_{e_a}\right)+\sum_{e \notin \mathcal{E}_A} \exp \left(\mathbf{x}_e^{(L) \top} \mathbf{u}_e\right)} \tag{9}
Lft?=?∣EA?∣1?ea?∈EA?∑?logexp(xea?(L)??uea??)+∑e∈/EA??exp(xe(L)??ue?)exp(xea?(L)??uea??)?(9)
即对每个可行答案的Softmax求平均。
其中，当计算一个答案的损失时，忽略其他答案的预测。 u e \mathbf{u}_e ue?是指实体 e e e在解码器中的嵌入。
在实践中，一种直观的策略是在每个推理任务的相应下游训练集上微调预训练的kgTransformer。然而，最近的NLP研究表明，多任务微调可以避免过拟合，让任务相互受益。受此研究启发，我们先在所有可能的下游训练集上联合微调kgTransformer，然后针对每一项任务再进行微调。此外，我们观察到，几个下游训练集的组合有时可以比单个任务的训练集好得多。最后，根据每个任务验证集的性能选择用于测试的最佳检查点。
3.2.3 域外泛化
为了更好地测试模型对未知查询的泛化能力，四种特定类型的域外查询 i p , p i , 2 u , u p ip,pi,2u,up ip,pi,2u,up仅在验证集和测试集中提供，如图2(b)所示。
其中，合取(∧)查询 i p ip ip和 p i pi pi可以表示为查询图，直接分解为结构子合取查询。对于析取(并集∨)查询 2 u 2u 2u和 u p up up，我们采用析取范式的思想(一个析取∨可以分解成多个合取∧)。具体来说，我们首先分别预测分解后的子合取查询，在分别求得子查询的结果后，将其概率转换为概率排名，最后取总排名第一作为最终的预测结果，（相当于求整个组合的最优解，最好不使用贪婪法来取概率最大值）。
图2(b)：掩蔽预训练与微调。