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

早前 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