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

©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，而应用到多跳的逻辑查询）。