[论文解读] DGL-KE: Training Knowledge Graph Embeddings at Scale

Da Zheng, Xiang Song+7 more|arXiv (Cornell University)|Apr 18, 2020

Advanced Graph Neural Networks参考文献 17被引用 36

一句话总结

DGL-KE 引入用于在大图上可扩展训练知识图嵌入的新优化，在多-GPU、众核 CPU 和分布式设置下相较于 GraphVite 和 Pytorch-BigGraph 取得显著加速。

ABSTRACT

Knowledge graphs have emerged as a key abstraction for organizing information in diverse domains and their embeddings are increasingly used to harness their information in various information retrieval and machine learning tasks. However, the ever growing size of knowledge graphs requires computationally efficient algorithms capable of scaling to graphs with millions of nodes and billions of edges. This paper presents DGL-KE, an open-source package to efficiently compute knowledge graph embeddings. DGL-KE introduces various novel optimizations that accelerate training on knowledge graphs with millions of nodes and billions of edges using multi-processing, multi-GPU, and distributed parallelism. These optimizations are designed to increase data locality, reduce communication overhead, overlap computations with memory accesses, and achieve high operation efficiency. Experiments on knowledge graphs consisting of over 86M nodes and 338M edges show that DGL-KE can compute embeddings in 100 minutes on an EC2 instance with 8 GPUs and 30 minutes on an EC2 cluster with 4 machines with 48 cores/machine. These results represent a 2x~5x speedup over the best competing approaches. DGL-KE is available on https://github.com/awslabs/dgl-ke.

研究动机与目标

• 推动在拥有数百万节点和数十亿边的巨大知识图上实现可扩展的嵌入训练。
• 开发并实现优化，以改善数据局部性、减少通信并在 KGE 训练中实现计算重叠。
• 提供一个基于 DGL、PyTorch/MXNet 以及分布式 KVStore 的开源包（DGL-KE），以实现大规模训练。

提出的方法

• 将知识图按计算单元划分以最小化跨分区数据移动。
• 基于 METIS 的图划分以将嵌入与三元组放在同一位置，减少跨机器通信。
• 实现联合负采样和基于度的负采样以减少内存搬运并生成更难的负样本。
• 对关系嵌入进行分区并传播稀疏读取/更新，以最小化大量关系的数据传输。
• 将梯度更新与批处理重叠，以隐藏 CPU-GPU 延迟并提高吞吐量。
• 引入周期性同步与分布式键值存储来协调跨机器的训练。

实验结果

研究问题

• RQ1 如何在具有数千万节点和数亿边的图上高效地进行大规模知识图嵌入训练？
• RQ2 哪种图划分、负采样和数据放置的组合在速度与嵌入质量之间能实现最佳折衷？
• RQ3 异步、多进程和分布式训练是否能够在降低数据传输和内存瓶颈的同时实现有竞争力的嵌入质量？
• RQ4 在大规模知识图上，DGL-KE 的性能与现有工具（GraphVite、PBG）相比如何？
• RQ5 关系分区和梯度更新重叠对训练吞吐量和模型准确性有何实际影响？

主要发现

• DGL-KE 实现了显著的加速，例如在一个 8-GPU EC2 实例上使用 100 分钟，在一个 4 机集群（每机 48 核）上使用 30 分钟，针对包含 8600 万节点、3.38 亿边的图，与竞争方法相比。
• 联合负采样在多-GPU 设置中将数据移动量降低高达约 40 倍，并提升张量运算效率。
• 基于度的负采样在大型图（如 Freebase）上提升模型准确性，尽管可能不会提高训练速度。
• CPU 与 GPU 之间的梯度更新重叠在大型图上对于许多模型可带来约 40% 的加速，同时保持准确性。
• 关系分区显著降低数据移动量，并在具有大量关系相关参数的模型（如 TransR）上带来显著的加速。
• 总体而言，DGL-KE 与 GPU 数量几乎线性扩展，能够在大型图上数小时内训练复杂的 KGE 模型，在公开实验中优于 GraphVite 和 PBG。