# DGL-KE: Training Knowledge Graph Embeddings at Scale

arXiv: https://arxiv.org/abs/2004.08532 · HTML: https://mlsanity.com/paper/2004.08532/

## Key facts

- Cited by 67 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 18 Apr 2020.

## Built on by (most cited first, 30 of 67)

- [OGB-LSC: A Large-Scale Challenge for Machine Learning on Graphs](https://mlsanity.com/paper/2103.09430/) (2021)
- [A Generalization of ViT/MLP-Mixer to Graphs](https://mlsanity.com/paper/2212.13350/) (2022)
- [Knowledge Graphs: Opportunities and Challenges](https://mlsanity.com/paper/2303.13948/) (2023)
- [A Survey on Graph Representation Learning Methods](https://mlsanity.com/paper/2204.01855/) (2022)
- [A*Net: A Scalable Path-based Reasoning Approach for Knowledge Graphs](https://mlsanity.com/paper/2206.04798/) (2022)
- [Marius: Learning Massive Graph Embeddings on a Single Machine](https://mlsanity.com/paper/2101.08358/) (2021)
- [MariusGNN: Resource-Efficient Out-of-Core Training of Graph Neural Networks](https://mlsanity.com/paper/2202.02365/) (2022)
- [Saga: A Platform for Continuous Construction and Serving of Knowledge At Scale](https://mlsanity.com/paper/2204.07309/) (2022)
- [SMORE: Knowledge Graph Completion and Multi-hop Reasoning in Massive Knowledge Graphs](https://mlsanity.com/paper/2110.14890/) (2021)
- [Context-aware Entity Typing in Knowledge Graphs](https://mlsanity.com/paper/2109.07990/) (2021)
- [Efficient and effective training of language and graph neural network models](https://mlsanity.com/paper/2206.10781/) (2022)
- [Friendly Neighbors: Contextualized Sequence-to-Sequence Link Prediction](https://mlsanity.com/paper/2305.13059/) (2023)
- [Distantly Supervised Relation Extraction with Sentence Reconstruction and Knowledge Base Priors](https://mlsanity.com/paper/2104.08225/) (2021)
- [Swift and Sure: Hardness-aware Contrastive Learning for Low-dimensional Knowledge Graph Embeddings](https://mlsanity.com/paper/2201.00565/) (2022)
- [NeuralKG: An Open Source Library for Diverse Representation Learning of Knowledge Graphs](https://mlsanity.com/paper/2202.12571/) (2022)
- [Explainable Biomedical Recommendations via Reinforcement Learning Reasoning on Knowledge Graphs](https://mlsanity.com/paper/2111.10625/) (2021)
- [Mining On Alzheimer's Diseases Related Knowledge Graph to Identity Potential AD-related Semantic Triples for Drug Repurposing](https://mlsanity.com/paper/2202.08712/) (2022)
- [Vehicle Lane Change Prediction based on Knowledge Graph Embeddings and Bayesian Inference](https://mlsanity.com/paper/2312.06336/) (2023)
- [GraphStorm: all-in-one graph machine learning framework for industry applications](https://mlsanity.com/paper/2406.06022/) (2024)
- [A Retrieve-and-Read Framework for Knowledge Graph Link Prediction](https://mlsanity.com/paper/2212.09724/) (2022)
- [PubGraph: A Large-Scale Scientific Knowledge Graph](https://mlsanity.com/paper/2302.02231/) (2023)
- [KGrEaT: A Framework to Evaluate Knowledge Graphs via Downstream Tasks](https://mlsanity.com/paper/2308.10537/) (2023)
- [Knowledge Graphs Querying](https://mlsanity.com/paper/2305.14485/) (2023)
- [PIE: a Parameter and Inference Efficient Solution for Large Scale Knowledge Graph Embedding Reasoning](https://mlsanity.com/paper/2204.13957/) (2022)
- [μKG: A Library for Multi-source Knowledge Graph Embeddings and Applications](https://mlsanity.com/paper/2207.11442/) (2022)
- [Tango: rethinking quantization for graph neural network training on GPUs](https://mlsanity.com/paper/2308.00890/) (2023)
- [The DLCC Node Classification Benchmark for Analyzing Knowledge Graph Embeddings](https://mlsanity.com/paper/2207.06014/) (2022)
- [Extreme Acceleration of Graph Neural Network-based Prediction Models for Quantum Chemistry](https://mlsanity.com/paper/2211.13853/) (2022)
- [Growing and Serving Large Open-domain Knowledge Graphs](https://mlsanity.com/paper/2305.09464/) (2023)
- [HDReason: Algorithm-Hardware Codesign for Hyperdimensional Knowledge Graph Reasoning](https://mlsanity.com/paper/2403.05763/) (2024)

## Abstract (arXiv)

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.