[论文解读] PyTorch-BigGraph: A Large-scale Graph Embedding System

Adam Lerer, Ledell Wu+5 more|arXiv (Cornell University)|Mar 28, 2019

Advanced Graph Neural Networks参考文献 40被引用 204

一句话总结

PyTorch-BigGraph (PBG) 将图嵌入扩展到十亿节点和万亿边，借助分区、分布式执行、批量负采样以及混合参数服务器方法，在性能竞争力的同时提升可扩展性。

ABSTRACT

Graph embedding methods produce unsupervised node features from graphs that can then be used for a variety of machine learning tasks. Modern graphs, particularly in industrial applications, contain billions of nodes and trillions of edges, which exceeds the capability of existing embedding systems. We present PyTorch-BigGraph (PBG), an embedding system that incorporates several modifications to traditional multi-relation embedding systems that allow it to scale to graphs with billions of nodes and trillions of edges. PBG uses graph partitioning to train arbitrarily large embeddings on either a single machine or in a distributed environment. We demonstrate comparable performance with existing embedding systems on common benchmarks, while allowing for scaling to arbitrarily large graphs and parallelization on multiple machines. We train and evaluate embeddings on several large social network graphs as well as the full Freebase dataset, which contains over 100 million nodes and 2 billion edges.

研究动机与目标

• 动机：对拥有十亿级节点和万亿条边的工业级规模图进行可扩展图嵌入的需求。
• 介绍 PyTorch-BigGraph (PBG) 及其可扩展性机制。
• 展示 PBG 在标准基准上的表现具有竞争力，同时能扩展到大规模图。
• 展示在 Freebase、LiveJournal、YouTube 和 Twitter 上的结果，以验证可扩展性和效率。
• 提供开源发布及针对大规模图嵌入管线的实用指南。

提出的方法

• 提出一个多关系图嵌入框架，使用分解的打分函数 f 来建模边及其关系特定的变换。
• 将邻接矩阵分成 P 个桶，在一次训练一个桶以适应内存约束并实现分布式执行。
• 利用基于锁服务器和分区服务器的并行化在多台机器之间协调训练，同时通过异步参数服务器共享少量全局参数。
• 实现内存高效的批量负采样，在一个批次内重复使用负样本，降低内存带宽。
• 使用 Adagrad 优化并采用基于边际的排序目标，支持可配置的采样与关系运算符（如 TransE、DistMult、ComplEx、RESCAL）。
• 支持单机和分布式训练，包括对分区实体类型与未分区实体类型的模块化处理，以及按关系的边配置。

实验结果

研究问题

• RQ1 能否在不牺牲嵌入质量的前提下，利用分区和分布式训练对拥有十亿级节点和万亿边的大规模图进行有效嵌入？
• RQ2 分区策略、负采样和关系运算符如何影响大规模真实世界图和知识图谱的可扩展性与模型精度？
• RQ3 在跨多台机器和分区进行 PBG 的训练时，训练时间、内存使用和收敛性存在什么权衡？
• RQ4 在不同评分运算符（TransE、ComplEx 等）下，PBG 在社会网络与知识图谱等多样化图类型上的表现如何？
• RQ5 将 PBG 方法用于发布大规模嵌入（如完整的 Freebase 图）是否具有实际可用性，且资源开销在可接受范围内？

主要发现

• PBG 在标准基准上与现有嵌入系统的性能相当，同时可扩展到拥有十亿级节点和万亿边的图。
• 分区大幅降低内存使用（例如对完整 Freebase 图的降低达 88%），嵌入质量仅有轻微下降。
• 在多台机器上的分布式执行可实现显著的加速（例如在 Freebase 上 8 台机器可实现高达 4x 的实际墙钟时间缩短），且在大图上接近线性扩展。
• 对负样本在一个批次内重复使用的批量友好型负采样显著提升训练速度，随着负样本数量增加仍保持高效。
• PBG 能在大型真实世界图（LiveJournal、YouTube、Twitter）上实现具有竞争性的链接预测与节点分类性能，嵌入仍可作为下游任务的特征使用。
• 在 Freebase 知识图谱上，PBG 搭配 TransE 与 ComplEx 达到与既有基线相当的平均逆排名（MRR）和 Hits@10；通过分区和分布式训练扩展到完整的 Freebase 图是可行的。