早前DeepMind开源了自己的图神经网络库Graph Nets library，在业界掀起了一股图神经网络的热潮。最近，Facebook开源了自己的图神经网络库PyTorch-BigGraph（简称PBG），有了它，再大的图都能快速生成图嵌入，并且完全不需要GPU，对拥有数十亿节点的大规模图嵌入、电子商务中的推荐系统、社交媒体中的链接预测等任务都提供了极大的便利。研究人员在最近的SysML会议上发表了一篇关于PBG的论文，文中公开可用的社交图数据集实验结果表明“PBG优于竞争方法”。本文是AI前线第75篇论文导读，我们将深入了解PBG背后的系统架构细节、原理和大型图上的实验效果。

介绍

图结构数据是各种机器学习任务常见的输入。直接处理图数据非常困难。一种常见的技术是使用图嵌入方法为每个节点创建向量表征，通过度量节点向量之间的距离来预测图中边的出现。图嵌入技术被证明是有用的。它为下游任务可以提供有用的特征，比如：电子商务中的推荐系统、社交媒体中的链接预测、预测药物相互作用等任务。

现在图数据在互联网公司中很常见，通常规模都比较大，这样就对标准的嵌入方法提出了额外的挑战。例如，Facebook图网络数据包含超过20亿用户节点和超过一万亿个喜欢、转发等关系的边。

这样大规模的图嵌入有两个主要的挑战，首先，嵌入系统必须足够快，以便在合理的时间内嵌入边为10的11次方到12次方的图。其次，具有20亿个节点和每个节点100个嵌入参数的模型（表示为浮点数）将需要800GB的内存来存储其参数，因此许多标准的嵌入方法已经超过了常用商品服务器的内存容量。

我们引入了PyTorch-BigGraph（下文简称PBG），这是一个嵌入系统，包括对标准模型的若干修改。PGB的作用是扩展到具有数十亿个节点和数万亿个边的图。PBG的重要组成部分有：

（1）邻接矩阵分块分解到n个桶中，每次从一个桶开始在边上训练。然后，PBG要么交换每个分区嵌入到磁盘（以减少内存使用），要么跨多台机器分布式执行。

（2）利用大型参数矩阵的块分解来分布式执行模型，以及全局参数和特征节点的特征嵌入的参数服务器体系结构。

（3）节点的高效负采样。从数据中均匀采样负样本节点，并在批处理中重用负样本节点，以减少内存带宽。

（4）支持多实体，多关系的图。它们可以配置相关选项，比如边权重、关系运算符的选择。
我们在Freebase、LiveJournal和YouTube图上评估了PBG，表明它与现有嵌入系统的性能相匹配。

我们还在较大的图上报告了实验结果。我们构建了一个完整的Freebase知识图（1.21亿个实体，24亿条边）的嵌入，并将其跟本论文一起公开发布。Freebase知识图的分区节省了88%的内存消耗（嵌入质量只有略微下降），8台机器上的分布式执行将训练时间减少了4倍。我们也在一个大的Twitter图上进行了实验，显示了近似线性变化的类似结果。PBG作为一个开源项目已经发布在GitHub开源社区上，它完全是用PyTorch编写，无外部依赖或者自定义运算符。

开源项目传送门：
https://github.com/facebookresearch/PyTorch-BigGraph

多重关系嵌入

3.1 模型

多重关系图是一个有向图G=(V,R,E)，其中V是节点（也可以称为实体），R是一组关系，E是一组边。通常元素e=(s,r,d)，其中，s为源，r为关系，d为目标。我们还讨论了具有多个实体类型的图。这些图有一组实体类型和从节点到实体类型的映射，并且每个关系为该关系的所有边的源节点和目标节点指定单个实体类型。

我们用参数向量表示每个实体和关系类型。我们将这个向量表示为θ。多重关系图嵌入使用一个分数函数 f(θs, θr, θd)，该函数为每一条边生成一个分数，该函数试图使任意(s,r,d)属于E的f(θs, θr, θd)的分数最大化，并将（s,r,d）不属于E的分数最小化。

PBG中一个经过转换的边的源和目标实体向量（θs, θd）之间的评分函数为：

其中，θr对应于特定关系转换运算的参数。使用因子分解评分函数生成一个嵌入，其中节点嵌入之间的（变换）相似性具有语义含义。

PBG使用点积或余弦相似度评分函数，选择关系运算g，包括线性变换、平移和复数乘法。评分函数和关系运算的组合允许PBG训练RESCAL、DistMult、TransE和ComplEx模型 (Nickel et al., 2011; Yang et al., 2014; Bordes et al., 2013; Trouillon et al., 2016）。关系类型的子集可以使用恒等关系，以便未转换的实体嵌入可以预测此关系的边。

我们考虑稀疏图，PBG的输入是一个带正标签（现有的）边的列表。负边样本是通过采样来构造的。在PBG中，通过在现存的边上采样一个新源或目标，来破坏正样本边以生成负样本。

由于现实世界图的边分布是重尾分布，因此选择如何采样节点来构造负样本可能会影响模型的质量（Mikolov et al., 2013）。一方面，如果我们严格按照数据分布对负样本进行抽样，模型对于预测罕见节点边的高分没有惩罚。另一方面，如果我们对负样本进行均匀采样，根据数据集中的源节点和目标节点频率对边进行评分，模型可以很好的运行（特别是在大型图中）。这两种结果都是不想要的，因此，在PBG中，根据它们在训练数据中的流行性我们采样α比率的负样本和随机采样(1−α)比率的负样本。默认情况下，PBG使用α = 0.5。

在多实体图中，负样本仅从边关系的正确实体类型中采样。因此，在我们的模型中，”无效“边（错误的实体类型）的分数是未定义的。以前在知识图的上下文中已经研究过使用实体类型的方法，但是我们发现在节点数高度不平衡的情况下，拥有实体类型的图中特别重要，例如，10亿用户vs100万产品。在所有节点上采用统一的负采样，损失将由用户负样本节点来控制，并且不会优化用户-产品边之间的排名。PBG优化了训练数据中每个边e和一组边e之间基于边缘（margin-based）的排序目标函数，这些边e是通过破坏（corrupt）被采样的源或目标节点（不能同时使用）来构造的。

其中，λ是边缘（margin）超参数和

为了重现某些图嵌入模型，还可以使用Logistic和softmax损失函数来代替排序目标函数。

模型通过小批量随机梯度（SGD）对嵌入和关系参数进行更新。我们使用Adagrad优化器，对每个嵌入向量的累积梯度G求和，以减少在大型图上的内存使用。

大规模训练

PBG设计用于操作任意大小的图，可以运行在单机上或者多台机器的分布式集群上。在这两种情况下，训练都发生在CPU线程数等于机器核数上，机器核之间没有显式同步。

4.1 实体和边的分区

PBG利用分区模式来支持较大的模型，这些模型通常无法在单机上装入内存。此分区还允许对模型进行分布式训练。

图G中的每个实体类型可以是分区的，也可以是未分区的。分区的实体被切分成P部分。P的选择使每部分都能装入内存，或者支持执行所需的并行度级别。

实体分割后，边根据其源实体和目标实体的分区划分到多个桶中。例如，如果边的源实体在分区p1中，目标实体在p2中，那么它将被放入桶（p1,p2）中。如果源和目标实体类型都被分区，那么将创建P的2次方个桶。如果仅仅是源实体（或者目标实体）被分区，那么将创建P个桶。

训练期间，每个epoch都迭代每个边桶。对于边桶（pi,pj），源分区i和目标分区j，分别从磁盘交换，然后，加载边（或者边子集）并在线程之间细分以便进行训练。

图分区针对上一节中描述的基本算法中引入两个函数更改。首先，在排序损失中，每个候选边（s,r,d）仅与负样本（s,r,d^）进行比较，其中，d^来自同一分区（与源节点相同）。

其次，边不再被采样，但是是按分区分组。在这种修改下仍然可以保证，SGD下收敛到一个平衡或者chain-recurrent点，但收敛速度较慢。

我们观察到遍历边桶的顺序可能会影响到最终的模型的效果。具体来说，对于除第一个之外的每个边桶（p1,p2），在前一次迭代中对边桶（p1,）或者（,p2）进行训练是重要的。此约束确保所有分区中的嵌入在同一空间中对齐。对于单机嵌入，我们发现，图1所示的”由内向外“排序，在最小化磁盘交换次数的同时获得了最佳性能。

4.2 分布式执行

现有的分布式嵌入系统通常使用参数服务器体系结构。在此体系结构中，参数服务器（可能是分片的）包含一个key-value存储的嵌入。在每次SGD迭代时，从参数服务器请求一个小批量数据所需的嵌入参数，并将梯度（异步）发送到服务器以更新参数。

参数服务器范式对于训练大型稀疏模型是有效的，但它有许多缺点。一个问题是基于参数服务器的嵌入框架需要太多的网络带宽才能有效运行，因为每个小批量边及其相关负样本的所有嵌入都必须在每个SGD步中传输。此外，我们发现有效的使用研究是必要的，同样的模型可以在单机或者分布式环境下运行，但是参数服务器架构限制了单机环境下运行模型的大小。最后，我们希望避免异步模型更新带来的潜在收敛问题，因为我们的嵌入已经被划分成独立的集合。

给定分区实体和边，PBG采用并行化方案，该方案将第4.1节中描述的模型分区上的锁模式与共享参数异步参数服务器体系结构相结合，即关系运算符、未分区或者特征化实体类型。

在这个并行化方案中，如图2所示，分区嵌入被机器在训练中锁定。多个边桶可以并行训练，只要它们在不相交的分区集上操作，如图1（左）所示。训练可以在P/2个机器上并行化进行。分区的锁定由一台机器上的集中锁服务来处理，为了最小化通信（即支持重新使用分区），锁服务器将桶分包给worker节点，还维护第4.1节中描述的不变量，即只有第一个bucket可以在两个未初始化的分区上运行。

分区嵌入本身存储在跨N个训练机器分片的分区服务器中。机器从分区服务器获取源和目标分区（通常大小为好几GB），并在从共享磁盘加载的边桶上训练，分区实体的检查点间歇性保存到共享磁盘。

有些模型参数是全局的，因此无法分区。最重要的是包括关系参数，以及基数非常小或使用特征化嵌入的实体类型。此类参数的数量相对较小（<106），它们通过一个分片参数服务器的异步更新进行处理。具体来说，每个训练机维护一个后台线程，可以访问所有未分区的模型参数。这些线程异步地从服务器获取参数并更新本地模型，并将累积的梯度从本地模型推送到参数服务器。该线程通过一些限制执行连续同步，以避免网络带宽饱和。

4.3 批量负采样

大多数图嵌入方法都使用负采样是跟内存（或网络）高度绑定，因为它需要B Bn d 浮点数次数的内存访问，来执行复杂度为O(B Bn d)的浮点操作。事实上，Wu等人报告说，训练速度”接近于逆线性函数(负样本数量)。

为了提高大型图的内存效率，我们观察到可以重复使用单批Bn采样源或目标节点来构造多个负样本。在一个典型的设置中，PBG从训练集中取一个批次B=1000的正边样本，并将其分成50个边的块。来自每个块的目标（同样，源）嵌入与从尾部实体类型均匀抽样的50个嵌入连接起来。50个正例与200个采样节点的外积等于9900个负样本。训练计算总结在图3中。

这种方法比在每批次中负采样开销低的多。对于每批次B个正边，仅从内存中提取3B个嵌入，并计算3 B Bn个边的得分（点积）。一个批次边的得分可以用一个批次的Bn x Bn矩阵乘法来计算，可以高效地执行。图4显示了不同数量的负采样的PBG的性能。

在关系数较少的多关系图中，我们构造了一批边，这些边都具有相关关系类型r。这提高了线性关系运算符fr(t)=Ar*t的训练速度，因为它可以被表示为矩阵乘法fr(T)=Ar T。

实验

5.1 实验设置

对于每个数据集，我们报告了learning rate为0.001-0.1、margin为0.05-0.2和negative batch size 为100-500的网格搜索的最佳结果。FB15k的结果在单独的测试集上报告。

所有实验都在具有24个Intel®Xeon®内核（两个插槽）和每个内核两个超线程的计算机上进行，总共48个虚拟内核和256 GB RAM。

我们用40 HOGWILD线程来训练。对于分布式执行，我们使用一组通过50GB/s以太网连接的集群机器。我们使用了torch.distributed的TCP后端，实际上实现了大约1GB/s的发送/接收带宽。对于内存使用测量，我们报告了以0.1秒间隔采样的峰值驻留集大小。

5.2 LiveJournal

我们对LiveJournal数据集上的PBG性能进行评估，在该数据集上，用户可以follow其他人形成社交网络。数据集包含4847571个节点和68993773个边。我们构建了包含75%边的训练集和包含25%边的测试集。

我们将PBG嵌入性能与MILE进行了比较，MILE也可以扩展到大型图。MILE反复地将图粗化为较少的图，并在每个级别的粗化图上应用传统的嵌入方法，最后进行细化，得到原始图的嵌入。文中还展示了基于MILE嵌入的DeepWalk的性能。

我们使用第5.4.2节中描述的相同方法评估嵌入。为了比较不同方法的计算时间，我们报告了训练期间不同方法获得的测试MRR的学习曲线（见图5）。

5.3 YouTube

为了证明PBG嵌入对下游监督任务有用，我们将PBG应用于YouTube数据集。数据集包含YouTube上用户之间的社交网络，以及这些用户的标签，这些标签表示他们订阅组的类别。此社交网络数据集包含1138499个节点和2990443个边。

我们通过将嵌入应用到用户多标签分类上，来比较了PBG、MILE和DeepWalk的嵌入性能。我们遵循典型的方法（Perozzi et al., 2014; Liang et al., 2018）来评估嵌入性能，这里我们通过随机选择90%的打标签数据作为训练集，剩余的10%作为测试集，进行10-fold交叉验证。我们以学习的嵌入为特征，训练一个one-vs-rest(一对多)的逻辑回归模型来解决多标签节点分类问题。
我们发现，PBG嵌入比竞争方法（见表1）的性能稍好。

5.4 Freebase Knowledge Graph

Freebase(FB)是一个包含从维基百科等提取的一般事实的大型知识图。FB15k数据集由FB的子集组成，包含14951个实体，1345个关系和592213个边。

5.4.1 FB15K

我们对基于PBG嵌入的链接预测任务和基于现存方法嵌入的知识图的性能进行了比较。我们对基于现存方法的知识图嵌入的MRR（mean reciprocal rank）和Top10命中（Hit@10）进行了比较。结果展示在表2中。

我们将FB15k与复杂的乘法关系运算符一起嵌入（Trouillon et al., 2016)。我们使用两种不同的配置来评估PBG：一种类似于TransE模型，一种类似于ComplEX模型。正如在这项工作中一样，我们也发现对为源负样本和目标负样本使用独立的关系嵌入是有益的(described as ‘reciprocal predicates’ in Lacroix et al. 2018)。对于ComplEX，我们在负样本上使用点积操作，利用softmax 损失，epoch为50训练400维的嵌入。

PBG的表现与TransE和ComplEx模型的报告结果相当。此外，最近的论文报告了FB15k（以及像WordNet这样的其他小型知识图）使用具有数千个维度的嵌入的ComplEx模型的获得非常好的结果（Lacroix等，2018）。我们设法在PBG框架中重现这些体系结构和结果，但由于篇幅，不在此处报告详细信息。

5.4.2 Full Freebase

接下来，我们使用完整的Freebase数据集比较不同数量的分区和分布式训练。我们使用完整数据集（至少出现5次的所有实体和关系），结果总共有121,216,723个节点，25,291个关系和2,725,070,599个边。我们切分构建训练集、验证集和测试集，分别包含总边的90%，5%，5%。我们使用的FB的全集的数据格式与第5.4.1节描述的freebase 15k数据集格式相同。

为了研究分区数的影响，我们将freebase节点均匀地划分为不同的分区数，并测量模型性能、训练时间和峰值内存使用情况。然后我们考虑在不同数量的机器上进行并行训练。对于机器数M，我们使用2M个分区（这是允许此并行级别的最小分区数）。请注意，完整模型大小（d=100）为48.5GB。

我们对每个模型迭代10个epoch进行训练，使用跟FB15k相同的网格搜索策略，对每个分区进行超参数搜索。对于多机的评估，我们使用跟单机中取得最佳性能一致的超参数设置。

我们使用与第5.4.1节所述类似的链路预测任务评估模型。然而，由于候选节点数量众多，对于eval集中的每一个边，我们根据它们在训练数据中的普遍性，从这些实体集合中选择10000个候选负样本节点，以产生负样本边，用它来计算MRR和Top10命中（hits@10）。

结果报告在表3中，训练时间和内存使用情况。

我们观察到，在单机上，峰值内存使用率随着分区数量几乎呈线性下降，但由于在I/O上花费了额外的时间，因此训练时间有所增加。在多机上，完整的模型是在机器之间而不是在磁盘上共享，因此2台机器的内存使用率更高，但随着机器数量的增加线性减少。训练时间也随着机器数量的增加而减少，尽管在多台机器上的训练也会有一些开销。这包括I/O开销和不完全占用的组合。出现占用问题是因为可能并非总是存在具有非锁定分区的可用存储桶以供机器工作。增加分区数量相对于机器数量将增加占用率，但我们不会详细检查这种权衡。

Freebase嵌入在经过10个epoch训练后，4台机器，无论有没有节点划分和并行化，其链接预测精度几乎相同。对于最高并行度（8台机器），MRR从0.171到0.163略有下降。

使用ComplEx模型训练的PBG嵌入在链路预测任务上比TransE表现更好，单个分区上，d=200时，实现MRR为0.24,Top10命中（Hits@10）为0.32。然而，我们的实验表明，训练具有多个分区和多台机器的ComplEx模型是不稳定的，并且在不同的副本中，MRR在0.15到0.22之间变化。通过PBG分区对ComplEx模型性能的进一步研究留待将来工作。

5.5 Twitter

最后，与5.4.2节中研究的Freebase知识图相比，我们考虑社会网络图上PBG的大规模。我们嵌入了一个公开的Twitter 12子图(Kwak et al., 2010) (Boldi & Vigna, 2004) (Boldi et al., 2011)，其中包含41652230个节点和1468365182个边的社交网络，具有一个”follow“的边关系。我们构建了训练集、验证集和测试集，分别包含90%、5%、5%的总边数。

在表4中，我们报告了10个训练周期后的MRR和Top10命中（Hits@10）,以及不同分区和并行模式的训练时间和峰值内存使用情况。结果与表3一致：我们观察到多台机器的训练时间减少，而没有损失链接预测精度的情况下。

在图7中，我们报告了在训练过程中使用不同数量的机器获得的测试MRR学习曲线。与图6中的Freebase知识库学习曲线相比，Twitter图显示了训练时间的线性比例，因为图是并行化分和训练的。

结论

在本文中，我们提出了一个嵌入系统PyTorch-BigGraph，它可以扩展为具有数十亿个节点和数万亿个边的图。PBG支持多实体、多关系图，并支持按关系配置，如边权重和选择关系运算符。为了节省内存使用并允许并行化，PBG将邻接矩阵分块分解为N个桶，一次从一个桶边上训练。

结果表明，用PBG训练的嵌入质量与现有嵌入系统的质量相当，训练时间短。Freebase 图的划分在不降低嵌入质量的前提下，可以减少80%的内存开销，而在8台机器上的分布式执行可以将训练速度提高4倍。我们的实验表明，嵌入质量对社交网络数据集的划分和并行化具有很强的鲁棒性，但在关系数大、度分布严重偏态、关系运算（ComplEx使用）等情况下，嵌入质量对并行化更为敏感。因此，改进这些更复杂模型的规模是未来研究的一个重要领域。

我们已经介绍了PBG在我们所知道的最大的公开图数据集上的性能。然而，PBG 体系结构的最大好处来自于比这些图大1-2个数量级的图，在这些图中需要更多细粒度分区，并显露出更多的并行性。我们希望这项工作和PBG的开源版本有助于推动更大图数据集的发布，并增加对更大图的研究和报告结果。

论文链接：https://arxiv.org/pdf/1903.12287.pdf

PBG开源项目地址：
https://github.com/facebookresearch/PyTorch-BigGraph

创作场景

完爆旧系统！Facebook 开源图神经网络库 PBG，无需 GPU 搞定数十亿节点图嵌入