WeKB on Plato：微信知识图谱表示学习实践

原创微信大数据团队微信后台团队

作为一种存储事实信息的结构，知识图谱近年来逐渐成为研究热点，图谱表示学习浓缩了KG信息，在自然语言处理与推荐等场景亦成为一种有效的信息辅助，锋芒初现。业界图谱数据日趋成熟，然而模型落地和业务应用却刚刚起步。为了提供更好的KG数据能力，微信图计算框架Plato与微信知识图谱WeKB共同研发了表示学习模型TransE+，在link prediction上相较最强baseline取得了3%+效果提升，在知识图谱支撑场景探索落地。

1 简介

知识图谱 (Knowledge Graph) 以图的形式承载了诸多客观世界知识。随着深度学习风行，在“万物皆可embedding”的思想下，KG领域的表示学习 (Knowledge Graph Embedding, KGE) 近年来蓬勃发展。以Trans*模型为代表的KGE算法作为提取KG信息的前置流程，在自然语言处理 (Natural Language Processing, NLP) 和推荐等诸多细分任务上有效注入了先验知识，显著提升模型效果。

图1. 知识图谱

KGE的效果在学界已形成了基本共识。我们期望借助KGE理论，在微信场景下深入挖掘大规模知识图谱数据，提供通用化的实体信息，为上层业务提供长远助力。本文介绍Plato联合WeKB团队对于知识图谱表示学习的落地探索。我们总结了学界模型的共性，在图引擎部分提供扩展支持。在经典模型TransE的基础上，我们实现升级版模型TransE+，通过修改模型训练策略，在link prediction任务上达到state-of-the-art效果。面对业务问题，我们也进行了模型初版改造与探索，期待拿到特化场景下的收益。业务与学术终归存在参差，如何实现可落地的模型创新，平衡性能与效果，答案还在路上。

2 学界现有工作

KG上的对象分为两种，即实体 (entity) 与关系 (relation, $r$ )，对应图上的节点与连边。对于有向图来说，实体又能进一步分为头实体 (head, $h$ )（对应有向图中的源点）与尾实体 (tail, $t$ )（对应有向图中的汇点）。三元组 $(h, r, t)$ 描述了KG中存储的一些客观世界的知识。而KGE的目标就是对于实体和关系都找到一种比较好的表达来得到一份浓缩后的KG。分析到这里，模型的优化目标已经初具雏形。我们需要基于KG图结构信息，找到一种度量 $(\vec{h}, \vec{r}, \vec{t})$ 的理想方式，以此来优化实体与关系的embedding。

TransE [1] 首次提出了一种 $(\vec{h}, \vec{r}, \vec{t})$ 的度量公式和完整的KGE训练流程，是KGE中翻译模型 (Translation Model)的鼻祖。TransE认为，对于三元组 (北京，国家，中国)来说，理想的embbedding关系为 $\vec{h} + \vec{r} = \vec{t}$ （如图1）。在这个逻辑下，KG上的拓扑关系能够直观反映在embedding中，例如所有满足国家为中国的城市应当具有相同的embedding。在上式的基础上定义距离公式为 $d(\vec{h}, \vec{r}, \vec{t})=\left \| \vec{h} + \vec{r} - \vec{t} \right \| _{2}^2$ ，损失函数写为pairwise hinge loss： $\sum_{(h,r,t)\in S}\sum_{(h',r',t')\in S'}\left [ \gamma + d(\vec{h}, \vec{r}, \vec{t}) - d(\vec{h'}, \vec{r'}, \vec{t'})\right ]_+$ 。训练时，拉取ground-truth的正例三元组集合，并通过一些随机策略生成负例三元组 ${(h', r', t')}$ ，优化上述损失，即可得到较为合理的实体与关系表达。

图2. TransE原理示意图

TransE模型物理含义清晰，但并非全无局限，具体体现在难以建模一对多/多对一/多对多关系与对称关系。以图3为例，前文中提到，在TransE的模型设定下，同为中国城市的北京与深圳应当具有相同的embedding。然而，深圳所属省份为广东，北京却并非如此。换言之，（北京/深圳，国家，中国）这个多对一关系与其它三元组形成了冲突，在训练时无法兼顾。图示中北京与深圳存在对称关系，即均为一线城市，需要同时满足 $e_{sz} + r = e_{pk}, e_{pk} + r = e_{sz}$ 时，情况关系embedding被优化为0，所有对称关系都无法很好地建模。

图3. TransE算法缺陷示意图

上述缺陷催生了一系列Trans*模型。比如TransH [2]提出可以将实体=embedding映射到另一个空间中建模。通俗说来，模型可以容许”北京“与”深圳“两个实体在原空间中embedding不同，只要映射后满足 $\vec{h}_\perp + \vec{r} = \vec{t}_\perp$ 即可。TransR [3]模型进一步提升模型复杂度和表达能力，认为针对不同的relation，应当使用不同的空间映射函数 $M_r$ 。其它Trans*模型针对实体间关系的稀疏性等问题，也针对性地设计了打分函数 $f(h, r, t)$ （详见表1）。

图4. Trans*算法横向对比

表1. 其它Trans*模型 [14]

除了Trans翻译模型系列，RESCAL、DisMult等语义匹配模型在计算匹配分数时考虑relation的信息，也能够较好地衡量三元组信息的合理性。学界KGE模型仍然创新不断，多元化的业务需求也会催化模型升级。Plato团队面向通用KGE场景，在现有图引擎能力的基础上提供KG采样接口，在PlatoDeep中实现TransE模型的升级版本TransE+，作为Trans*模型的试水初探。

3 TransE+ on Plato

3.1 架构：面向KG场景升级图引擎

图5. PlatoDeep架构

PlatoDeep是微信自研的大规模图计算框架，最上层的PlatoDeep结合TensorFlow生态，抽象GNN与传统图计算依赖的底层能力，提供采样点/邻居、加载点/边属性、随机游走等接口，解决工业规模图数据加载与采样的问题，同时为用户开放定制模型的能力。KGE学习的流程与GNN模型的差异主要体现如下：

表2. GNN模型与KGE模型差异对比

特别地，KGE中采样三元组的概念与GNN中采样边有所区别。在普遍意义上的图中，节点间连边默认为0/1条，无论是稀疏/稠密属性，在逻辑上皆为定长。而KG中有所区别，实体间关系可能不止一种，且关系(relation)形成的集合 $S_R$ 在扩张的KG中一定会增多。若将节点间关系类别泛化为边上属性，既无法解决增量任务中关系种类变多的问题，采样边的概率分布也与采样三元组有天然差异，可能影响到模型。此外，在推荐场景下，数据上游的用户点击/交互数据反馈形成的类型集合 $S_I$ 亦与知识图谱关系集合 $S_R$ 类似。为了更好地服务这些场景，我们对节点间关系类型进行了单独抽象，支持带权加载，可进行<节点，关系，节点>形式的采样。业务场景中常见的负例生成逻辑也都可满足。预处理负例可以与正例一起，通过节点相同、关系不同的两张图完成加载。若要运行时完成负例生成，可以进行两次随机点采样/进行batch内id shuffle完成全局负例采样/batch负例采样。如此一来，KGE模型依赖的图引擎操作便完备了。

3.2 模型

3.2.1 经典TransE：美中不足？

经典TransE的训练流程遵循标准的pair-wise训练：在全局中采样正样本三元组，通过随机替换head/tail，构造等量的负样本，与正样本计算hinge loss，从而优化embedding向量。

图6. TransE算法流程 [1]

这个做法有什么缺陷？我们先来看真实场景的KG子集的数据：（FB15K/WikiData5m分别采样自FreeBase/WikiData）

表3. 知识图谱数据集概览

从这份数据可以看出来，潜在连边数量在真实世界图中一般是远超真实连边的。也就是说，若要将负样本信息完全吸收到模型当中，势必需要多看一定倍数的负样本，不契合pair-wise hinge loss这种将正负样本一一映照的设定，也不利于做sample reweighting。事实上，学界在TransE之后的一些工作也通过实验证实TransE原文的效果有待提升。我们广泛调研了后来者的工作，结合自己的理解，对经典TransE进行了调整，以求模型不落下风。

3.2.2 升级TransE：State-of-the-Art

在learning-to-rank模型当中，point-wise作为一种有效的模型训练方式，截止今日依旧在许多SOTA模型中使用。point-wise是将正负样本全部当做单条分类问题处理，试图直接得到样本正负性，即

$\text{argmin}_\Theta \mathcal{L}\left ( y, p||\Theta \right ) = \text{argmin}_\Theta \sum_{(h^{i},r^{i},t^{i},y^{i})\in \mathcal{D}}\alpha^{i}\mathcal{l}\left ( y^{i}, p^{i}||\Theta \right )$ 其中， $\mathcal{D}$ 表示全量正负样本集合， $\alpha^{i}$ 为样本权重， $p^{i}$ 是模型基于当前 $(h^{i}, r^{i}, t^{i})$ 计算得到的正负概率，单条样本损失函数 $l$ 可以选择交叉熵。采用这种训练方式的好处有二：一来可以使模型在无需重复学正样本的情况下看到更高比例的负样本，二来我们能够引入一些策略调整 $\alpha^{i}$ ，促使模型在训练后期能够使用Curriculum learning (CL) 策略，着重对未预测好的样本学习。 $p^i$ 与 $d(\vec{h}, \vec{r}, \vec{t})$ 呈负相关，即距离越大越趋负，

$p^{i}=\text{sigmoid}(\gamma -d(\vec{h}, \vec{r}, \vec{t}))$ $d(\vec{h}, \vec{r}, \vec{t})=\sum_k\left | h_i + r_i - t_i \right |$ $d(\vec{h}, \vec{r}, \vec{t})$ 我们在这里采用一阶范数计算， $k$ 指代embedding维度大小， $e_i$ 是 $\vec{e}$ 第 $i$ 位元素。向量初始随机化采用区间均匀分布策略。特别地，为了让embedding初始时能够落在合理的取值区间，我们设定区间为 $\left [ -\frac{\gamma}{k}, \frac{\gamma}{k} \right ]$ (margin based initialization)。由于正样本在训练时反复喂入，训练较为充分，我们着重对负样本计算差异化权重 (self-adversarial negative sampling)，策略如下，

$\hat{\alpha }_{\text{neg}}^{i} = \text{exp}((p_{\text{neg}}^i-\hat{p}_{\text{neg}}^i)/\theta _{\text{cur}})$ $\alpha _{\text{neg}}^{i} = \frac{\hat{\alpha }_{\text{neg}}^{i}}{\sum_{i\in \mathcal{B_{\text{neg}}}} \hat{\alpha }_{\text{neg}}^{i}}$

其中 $\hat{p}_{\text{neg}}^i$ 为样本i所在批次 (batch)中计算出的负样本的平均概率。 $\theta _{\text{cur}}$ 可视作这个CL策略中的松弛因子，越大越平滑。我们训练优化后的模型TransE+，在知识图谱补全任务上进行了验证，和原paper与学界的复现版本对比如下（数据集：FB15K）。

表4. Plato模型效果对比

对比下可以看出，TransE+ on Plato效果有较显著的优势。我们也针对上述优化进行了消解实验，对比[w/o xxx]与[TransE+ on Plato]可以清楚量化每一种策略带来的效果提升。w/o cross entropy是用经典TransE的hinge loss替代point-wise且采用1:1负采样，w/o margin based initialization使用[-1, 1]区间均匀分布，w/o self-adversarial negative sampling使用同样的负样本权重。

表5. Plato模型消解实验

从实验中我们能够发现，point-wise/margin based initialization对效果提升都十分显著：point-wise的效果增益主要来源于更丰富的负样本，margin based initialization帮助模型跨越到更小范围的embedding优化中。self-adversarial negative sampling策略下的模型在MR上效果稍差，但在MRR与Hits指标上拿到了显著的增益。

3.2.3 面向业务问题：模型剪枝与泛化

在学界，知识图谱实体表示学习一般应用于一些需要强通用/专业事实的场景，例如知识问答 (Knowledge Base Question Answering, KBQA)、开放式/任务型对话 (Open Domain/Task Oriented Dialog Generation)等。针对给定的简单问句，分段式的KBQA [8,9,10] 拆分实体链接 (Entity Linking) 和关系识别 (Relation Detaction)两个子任务，显式抽取或用模型编码实体和关系，再在图谱中查询或匹配得到答案。一些预训练语言模型，如THU-ERNIE [11]，也探索使用上下文信息融合KG知识来得到更好的embedding辅助下游任务。这些成果都证明了KGE在NLP任务中的有效性。

图7. KBQA流程示意图

理想很丰满，现实很骨感。学界普遍解决的是定义好的标准问题，在信息完备的数据集上实验，而我们在业界面临的问题却可能与学界存在些差异和冲突。以视频号推荐场景为例，视频tag匹配时没有上下文信息指代relation，亦无需上复杂模型完成推理，线上仅做相似度检索。这个问题与标准KGE的冲突在于，KGE建模 $f(h, r, t)$ ，在TransE模型中 $f(h, r, t)=h+r-t=0$ ；而相似度检索建模 $f(h, t)$ ，计算cosine相似度期望 $h-t=0$ 。以图2所示为例，在TransE模型中能够满足embedding近似的只有两类节点，一是与同样的头/尾实体有相同关系连接的节点（如北京、深圳均为中国城市），二是形成对称关系的节点（relation信息近似优化为0），KG上的一阶连接信息很难在similarity search中被检索出来。我们在WeKB上运行TransE+，对比两种验证方式。一种是采用学界标准的验证做法，即替换三元组head/tail计算link prediction效果。另一种与业务场景对齐，不考虑relation向量，仅使用head/tail检索相似实体。

表6. WeKB效果实验

注：WeKB数据集基于开源中文知识图谱构建，与用户个人数据无关。

实验结果也印证了前述的分析，即标准KGE直接用于相似度检索会带来大幅效果损失。我们从问题出发，思考模型改造的方向。既然场景中的确不需要关系信息，那么我们便回退到similarity search的场景下，裁剪模型中考虑relation信息的部分，也忽略KG边的有向性，采用双塔模型去优化实体表示即可。在KG构造时，实体描述携带着一些先验的实体语义信息。我们将这部分信息引入到模型中，弥补KG信息的不足，提升模型的泛化性。

图8. 面向业务问题的模型

在实体信息建模时，我们依旧做lookup来获取实体节点的结构embedding信息 $h^g/t^g$ 。实体描述信息使用一次单层映射作为语义域的隐式表示 $h^s/t^s$ 。为了消除结构域与语义域的概率空间差异，我们计算对称KL散度损失，

$\mathcal{L}_{\text{KL}}(e^g,e^s)=\frac{1}{2}( \sum _k e^g_i\text{log}\frac{e^g_i}{e^s_i} +\sum _k e^s_i\text{log}\frac{e^s_i}{e^g_i})$ 计算头尾实体相似度简化如下，其余细节与3.2.2中的模型无异。

$d(\vec{h}, \vec{r}, \vec{t})=\sum_k\left | h_i - t_i \right |$ 我们训练改进后的模型，在WeKB上采用实体similarity search来做embedding验证，相较于3.2.2所述TransE+模型，对齐了任务训练目标，在link prediction任务上MR指标也从40017提升至10399，模型泛化性提升显著。在业务实际使用过程中，我们也期望能够拿到相较于Graph Embedding (DeepWalk) 等方法更好的收益。

3.3 性能：多多并发起来

落地KGE的另一个问题是性能。知识图谱天然变动较小，不要求实时训练推理能力，但在离线场景下依然可以追求尽快完成高质量训练。我们对训练profile观察发现，图采样未对训练构成瓶颈。而由于模型训练部分在TTensorFlow分布式架构完成，所以主要的思考点是如何利用TF能力来有效优化训练性能。在异步训练的大框架下（图9），我们尽可能在单机上提升算子间/算子内并行度，提高worker的cpu使用，让batch计算变快。在异步更新框架下，太过激进地扩充worker虽然可以让训练速率看似变快，却可能因梯度失效影响模型效果。所幸在比较简单的KGE模型中（如3.3.2中所述的TransE），模型参数只有embedding lookup table。一般来说，KG的三元组涉及的实体节点和关系总数够大，且边关系比较均匀。以点规模400w、batch size 2000为例，worker同时命中相同节点的概率实际上很低。因此对于TransE，不需要对worker数量有太多顾忌。训练有共享参数的模型时（如3.3.3中所述的模型），再将worker的数量降低。

图9. TensorFlow分布式训练 [12]

除了调整TensorFlow并行度，我们对模型细节也进行了分析优化，包括在embedding lookup前后使用unique+gather策略减少重复lookup，替换了一些低效算子(map_fn)。上述策略的训练速率提升倍数量化如下：

表7. 训练性能优化

（注：上述策略实验控制了策略调整前后其它变量，但最终速率增长倍数并非累乘）

在节点数量421w、训练集数量1815w情形下，训练851798步（约合400个epoch）耗时123分钟，平均训练速率为130~150 global_step/sec。

4 总结与展望

本文从学界KGE研究成果出发，介绍了我们在Plato上研发TransE+的模型创新与工程细节，探索了微信知识图谱WeKB的KGE落地。但探索KG信息远不止于此，本文未尽工作不完全囊括如下：

研究更好的负采样策略：我们在模型研发的过程中尝试了一些负样本生成策略，包括使用全局带权负采样替代batch内shuffle负采样、先验计算hard negative samples等，然而并未拿到收益，还带来了额外的模型预处理和训练时开销。采用local closed world assumption [13]（假设KB信息是局部完备的，有选择地生成负例）等其它负样本生成方法依然留待实验验证。
端到端业务模型创新：下游业务逻辑带来的新增上下文和constraint都为模型改进带来了广阔空间，这一点还需要和业务同学共同交流创新。

参考文献

[1] Bordes, Antoine, et al. "Translating embeddings for modeling multi-relational data." Advances in neural information processing systems 26 (2013).

[2] Wang, Zhen, et al. "Knowledge graph embedding by translating on hyperplanes." Proceedings of the AAAI conference on artificial intelligence. Vol. 28. No. 1. 2014.

[3] Lin, Yankai, et al. "Learning entity and relation embeddings for knowledge graph completion." Twenty-ninth AAAI conference on artificial intelligence. 2015

[4] https://github.com/thunlp/OpenKE/tree/OpenKE-Tensorflow1.0

[5] Nickel, Maximilian, Lorenzo Rosasco, and Tomaso Poggio. "Holographic embeddings of knowledge graphs." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 30. No. 1. 2016.

[6] https://graphvite.io/docs/latest/benchmark.html#knowledge-graph-embedding

[7] Wang, Xiaozhi, et al. "KEPLER: A unified model for knowledge embedding and pre-trained language representation." Transactions of the Association for Computational Linguistics 9 (2021): 176-194.

[8] Yu, Mo, et al. "Improved Neural Relation Detection for Knowledge Base Question Answering." Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2017.

[9] Huang, Xiao, et al. "Knowledge graph embedding based question answering." Proceedings of the twelfth ACM international conference on web search and data mining. 2019.

[10] Luo, Kangqi, et al. "Knowledge base question answering via encoding of complex query graphs." Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing. 2018.

[11] Zhang, Zhengyan, et al. "ERNIE: Enhanced Language Representation with Informative Entities." Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics. 2019.

[12] Abadi, Martín, et al. "TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems." (2015).

[13] Dong, Xin, et al. "Knowledge vault: A web-scale approach to probabilistic knowledge fusion." Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining. 2014.

[14] Ji, Shaoxiong, et al. "A survey on knowledge graphs: Representation, acquisition, and applications." IEEE transactions on neural networks and learning systems 33.2 (2021): 494-514.

继续滑动看下一个