DGL-KE: Training Knowledge Graph Embeddings at Scale

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.

References in corpus (5)

• MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
• Graph Embedding Techniques, Applications, and Performance: A Survey
• Complex Embeddings for Simple Link Prediction
• PyTorch-BigGraph: A Large-scale Graph Embedding System
• Analysis of the Impact of Negative Sampling on Link Prediction in Knowledge Graphs

Cited by in corpus (14)

• Open Graph Benchmark: Datasets for Machine Learning on Graphs
• Drug Repurposing for COVID-19 via Knowledge Graph Completion
• PyKEEN 1.0: A Python Library for Training and Evaluating Knowledge Graph Embeddings
• Mining On Alzheimer's Diseases Related Knowledge Graph to Identity Potential AD-related Semantic Triples for Drug Repurposing
• Few-shot link prediction via graph neural networks for Covid-19 drug-repurposing
• DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphs
• CoLAKE: Contextualized Language and Knowledge Embedding
• Marius: Learning Massive Graph Embeddings on a Single Machine
• A Distributed Multi-GPU System for Large-Scale Node Embedding at Tencent
• Repurpose Open Data to Discover Therapeutics for COVID-19 using Deep Learning
• SMORE: Knowledge Graph Completion and Multi-hop Reasoning in Massive Knowledge Graphs
• CS-lol: a Dataset of Viewer Comment with Scene in E-sports Live-streaming
• PanRep: Graph neural networks for extracting universal node embeddings in heterogeneous graphs
• Scalable and interpretable rule-based link prediction for large heterogeneous knowledge graphs