KDD 2022 | kgTransformer：基于知识图谱与Transformer的复杂逻辑查询

科技

2022-09-05 02:09

©PaperWeekly 原创 · 作者 | 张成蹊

单位 | Freewheel机器学习工程师

研究方向 | 自然语言处理

Preface

最近在 KDD'22 上看到了一篇用 transformers 结构来做 KGe（Knowledge Graph embedding）的 paper，体裁比较新颖。

论文标题：

Mask and Reason: Pre-Training Knowledge Graph Transformers for Complex Logical Queries

论文链接：

https://arxiv.org/abs/2208.07638

代码链接：

https://github.com/THUDM/kgTransformer

论文有几方面的亮点：

1. 把 KG 里面的节点关系（relation/edge）转换成了关系节点，从而把两个节点+两者的关系变成了一个三元组。文中称之为 Triple Transformation Method；

2. 使用了两阶段 pre-training。第一阶段主要是通过两种 Random Walk 来对整个图进行随机采样，然后进行训练；第二阶段是通过预定义的一些范式（1p/2p/3p/2i/3i）来采样并训练；文中分别称为 Dense initialization 与 Sparse refinement；预训练的任务与 BERT 相类似，都是通过 mask 其中若干节点，并要求模型预测这些节点；（值得注意的是，在二阶段 pre-training 与整个 fine-tune 过程中，也会 mask 若干节点，但只要求模型预测出最终节点）；

3. 在 Multi-Head Attention 上层训练多个 FFN，以模拟出多专家投票的效果，来增大模型参数量；使用 Gating 的方式控制每次只选择 2 个 FFN 产出结果，以节省计算时间。文中称之为 Mix-of-Experts（MoE）。

文中提出的这些方法使得 kgTransformer 具有以下优势：

1. 动态的 KG embedding（因为使用了 transformer，此处可以参照 GloVe embed 与 ELMo embed）；

2. 模型的学习过程与应用的目标是匹配的（这句话是针对部分模型，训练目标是补全一阶的 graph relation link，而应用到多跳的逻辑查询）。

Introduction

2.1 Complex Logical Query

本文主要解决的问题是复杂逻辑查询，包括：补充边（imputed edges）、多源的实体、Existential Positive First-Order（EPFO）逻辑，以及未知的中间实体等。下面是 EPFO 的一个示例：

具体地，对于一句自然语言：

What musical instrucments did Minnesota-born Nobel Prize winners play?

所解析得到的 EPFO 查询，里面蕴含了一个若干约束的二跳（2-hop）知识图谱查询：

1. 第一个约束是查询一个实体，出生在 Minnesota；

2. 第二个约束是查询一个实体，获得了诺贝尔奖，与上一个约束指向同一个实体；

3. 第二跳是查询上面得到的实体所 Play 的乐器。

可知，第二跳的查询结果才是问句所想要得到的答案。

此类 EPFO 查询的难点在于：

1. EPFO 查询的解析复杂度与 hop 之间的关系是指数级的；

2. EPFO 查询要求模型具有比较强的迁移与泛化能力。

2.2 Contributions

论文的贡献在于：

1. kgTransformer：提出了 kgTransformer，使用 Triple Transformation 策略来将知识图谱带有属性的边（relation）转换为不带属性的有向边；提出了 MoE 来顺应 GNN 稀疏激活的特性。

2. Masked Pre-Training：为了提升泛化能力，作者将复杂逻辑查询问题看作一个预测 mask 的问题，通过两阶段预训练来训练模型；在 FB15k-237 与 NELL995 两个数据集上进行了验证，说明了模型具有 SoTA 的效果。

KG Pre-Training Framework

3.1 kgTransformer

kgTransformer 主要思想是把图作为 sequence，放到 transformer 里面进行预训练与微调。为了实现这个目的，首先需要进行数据的准备。

怎么将图转换为序列呢？节点之间的邻接关系可以通过 random walk 的方式获得，但 relation edge 带有的属性不好嵌入到这样的结构中去。所以作者提出在做数据准备的时候，把所有的 relation edge 也转换成节点加入的图中，而由 relation node 所连接的两个 entity node，所使用的边都是有向边。

值得注意的是，在这样的图里，实体节点与这个关系节点是需要做类型区分的，作者在 transformer 的输入层加入一个特殊的 node type embedding 来实现了这个目的。

文中所使用的 transformer 大体结构及其内部公式与 Attention Is All You Need 是一致的，只是 sequence 中每个 token 从原来的文字变成了现在的节点。在 encoding 的时候，通过 Multi-Head Attention 来计算得到每个节点相邻节点与它的关系重要程度，并通过 FFN 来进行变换得到最终输出。

有一点与 transformer 不一样的是，对于 MHA 上面的 FFN，作者初始化了多达 32 个（实际实验中选择的数量）FFN 来进行学习，以模型专家投票的机制，称为 Mix-of-Expert（MoE）。

在训练时，门控机制只会选择权重最大的两个 FFN，将其输出进行求和作为最终输出，以保证训练过程中的时间不过长；
在预测时，不使用门控机制，将所有的 FFN 结果求和作为最终输出。

为什么加入这样的 MoE 会让效果变好呢？作者提出：在 transformer中，已经有研究发现 FFN 与 key-value network 效果是等同的，即 FFN 通过全连接的方式实现了与 attention 类似的效果，其全连接的权重可以从侧面被看做是输入间的重要性关系。因为这个原因，transformer 里 FFN 输出的元素里大部分都是 0，是很稀疏的。

在此基础上，作者也做了一些前置研究，发现 EPFO 里的 FFN 一般只有 10%-20% 的神经元被激活（除了最后的 decoder 层）。正因为有稀疏性的存在，作者认为使用 MoE，把一个很大的 FFN 分解成若干个小的 FFN，能够更好的增强模型泛化能力。

3.2 Masked Pre-Training and Fine-Tuning

在定义 kgTransformer 的结构与输入数据模式之后，开始进行模型的预训练与微调。作者又把预训练分为了两个阶段，在第一个阶段，使用随机游走采样出不同的子图，进行模型的初始化；在第二个阶段，使用预先定义好的一些 meta-graph 结构继续预训练。

为什么要分成两个阶段进行预训练呢？作者认为，第一个阶段是尽可能将 KG 中的知识嵌入到模型中去，在此时采样窗口一般会设置得比较大，采样结果是稠密（dense）的，并且追求采样出来的子图的多样性，从而让模型见识尽可能多的 pattern；但是，在实际应用时，大部分的逻辑查询结果是稀疏（甚至 out-of-domain）的，而只进行第一阶段的预训练，会让模型的迁移能力变差。

3.2.1 2-stage pre-training

在预训练的第一个阶段：dense initialization，作者使用了两类随机游走的策略对原始图像进行采样，分别是 random walk with restart (RWR) 与 tree-based RWR。专门加入 tree-base 是因为 EPFO 的逻辑连接关系更像是树状的结构。

设 kgTransformer 内部参数为，随机 mask 的实体集合为，随机采样后的子图原始输入为，mask 之后的子图输入为，第一阶段 dense initialization 的学习目标是最小化损失函数：

其中表示节点之间的邻接关系，是由随机游走采样得到的。

在预训练的第二个阶段：sparse refinement，作者使用了基础的 EPFO pattern: 作为 meta-graph，保证每个子图中不会包含超过 4个实体，随后仍然采用 mask 的方式让模型预测 mask 前的 entity embedding，但在这个阶段，模型只会最小化所有目标实体的损失函数，即：

3.2.2 fine-tuning

尽管预训练的第二阶段已经比较接近于想要解决的任务，但作者发现使用目标数据集进行 fine-tune 仍然是比较重要的（原文附录表 2 中进行了论证）；此处的训练目标是最小化所有可行解的损失，给定一组可行的答案，微调阶段对于每个查询的损失函数记为：

其中加入的意为这是 softmax 的最终输出，整个式子要求所有可行实体的 softmax 输出概率之和尽可能大。

此外，作者注意到最近的 NLP 研究指出多任务学习能够有效避免模型过拟合，且从彼此的训练过程中受益。因此作者首先在所有的任务中进行了预训练，随后逐个在不同的任务中进行了微调。而在 out-of-domain 的泛化中，部分 pattern 是由几个 meta-graph 查询组合得到的，对于这些特殊的 pattern，作者采用了一个比较 tricky 的分解预测的方法，分别求得子查询的结果后，将其概率转换为概率排名，取第一名作为最终的预测结果。

Experiments

作者在 FB15k-237 与 NELL995 两个数据集的 9 个 reasoning 任务（包含了 in-domain 与 out-of-domain）中测试了 kgTransformer 的泛化能力，结果如图所示：

其中加粗部分为所有实验中效果最好的，下划线部分为实验中效果次好的。可以发现，与此前的 SoTA 方法 CQD 相比，作者提出的 kgTransformer 在 NELL995 上取得了绝对值 2.4%（相对值 4.6%）的提升，在 FB15k-237 数据上取得了绝对值 3.5%（相对值 12.1%）的提升，提升效果比较显著。

作者对 kgTransformer 的组成方法：2-stage pre-training 与 fine-tune 进行了缺省实验，结果如下所示：

作者发现，pre-training 与 fine-tuning 作为模型的训练过程都是十分重要的，不管缺少哪一个部分，都会带来非常大的效果下降。而对于 pre-training 过程中的两个阶段，其中第二个阶段对于 FB15k-237 几乎没有影响，而给 NELL995 带来了显著的提升。考虑到 sparse refinement 的提出动机，说明 NELL995 中的稀疏 pattern 更多。

对于 MoE 的效果，论文进行了不同数量 MoE 下的性能与模型效果比较，证明了使用合适数量（实验中为 32 个）的 Experts+Gating 机制，能够给 kgTransformer 带来 16 倍的 size 提升，而实际的计算成本只增加了 11.6%-38.7%。

更多阅读