Source-linked AI summary
DGL-KE: Training Knowledge Graph Embeddings at Scale
Da Zheng, Xiang Song, Chao Ma, Zeyuan Tan, Zihao Ye, Jin Dong, Hao Xiong, Zheng Zhang, George Karypis
TL;DR
Large knowledge graphs make embedding computation expensive because existing parallel approaches incur substantial data-transfer overhead and low computational efficiency. DGL-KE introduces locality, communication, parallelism, and asynchronous-update optimizations for multi-processing, multi-GPU, and distributed training. On a graph with over 86M nodes and 338M edges, it computes embeddings in 100 minutes with 8 GPUs or 30 minutes on four 48-core machines, outperforming competing packages in training time while maintaining comparable quality.
Problem
Growing knowledge graphs require computationally efficient embedding algorithms because existing parallel approaches suffer from high data-transfer overhead and low computational efficiency.
Method
DGL-KE combines data-locality, graph-partitioning, negative-sampling, asynchronous-update, and hardware-parallelism optimizations in an open-source KGE package.
Results
DGL-KE computes embeddings for a graph with over 86M nodes and 338M edges in 100 minutes with 8 GPUs or 30 minutes on four machines with 48 cores/machine, with comparable quality to competing approaches.
Takeaways & Limitations
DGL-KE significantly outperforms state-of-the-art KGE packages across many-core CPU, multi-GPU, and machine-cluster hardware while maintaining high model accuracy.
Takeaways & Limitations
Relation partitioning restricts which relations appear in a mini-batch and can reduce stochastic-gradient-descent randomization, potentially affecting embedding quality.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
DGL-KE addresses the high data-transfer overhead and low computational efficiency of existing approaches for large knowledge graphs with parallel training optimizations. It achieves comparable embedding quality while substantially reducing training time on very large graphs.
- Existing GraphVite and Pytorch-BigGraph approaches accelerate KGE with multi-GPU or distributed training but suffer from high data-transfer overhead and low computational efficiency.
- DGL-KE accelerates KGE training on graphs with millions of nodes and billions of edges through multi-processing, multi-GPU, and distributed parallelism.Its optimizations target data locality, communication overhead, computation-memory overlap, and operation efficiency.
- Its computation decomposition across cores, GPUs, and machines enables massive parallelization while reducing write conflicts and communication overhead.
- DGL-KE packages these optimizations as an open-source Python system built on DGL, with deep learning frameworks for tensor computation and a distributed key-value store.
- 100 minutes on an EC2 instance with 8 GPUs and 30 minutes on an EC2 instance with 4 machines with 48 cores/machine were required for a graph with over 86M nodes and 338M edges.These times correspond to embedding computation, with quality comparable to competing approaches.
- 5× and 2× speedups over GraphVite and Pytorch-BigGraph, respectively, were achieved on the reported large knowledge graph.
2 BACKGROUND
This section defines knowledge graphs, their embeddings, supported scoring models, training losses, and negative-sampling procedures. These components specify how entities and relations are represented and optimized.
- A knowledge graph represents typed entities and relations, commonly as triplets (h,r,t) connecting head entity h and tail entity t through relation r.
- Knowledge graph embeddings map entities and relations to low-dimensional vectors used in tasks such as knowledge graph completion and recommendation.The embedding dimension is denoted by d.
- KGE models optimize score functions to assign higher scores to observed triplets and lower scores to unobserved triplets.
- Supported models include translational distance models such as TransE and TransR and semantic matching models such as DistMult, ComplEx, and RESCAL.
- TransR is d times more computationally expensive than TransE because it applies additional matrix multiplications to head and tail embeddings.
- Training uses logistic or pairwise ranking losses over positive and negative triplets, with D+ and D− denoting the two sets and y denoting triplet labels.
- DGL-KE generates negative samples by corrupting entity heads or tails, while relation corruption is not used in this work.
3 METHODS
DGL-KE’s methods reduce data movement and update conflicts while exploiting CPUs, GPUs, and distributed machines. The system combines graph partitioning, negative sampling, hardware-specific data placement, and asynchronous updates.
- Motivation: Naive KGE training has computational density O(1), making data movement reduction essential because memory accesses are slower than computation.
- Motivation: Parallel training must exploit multi-core CPUs, GPUs, and machine clusters while minimizing conflicts during concurrent embedding updates.
- Hardware configurations: DGL-KE provides unified multiprocessing implementations for many-core CPU, multi-GPU, and CPU/GPU cluster configurations.
- Training pipeline: Training partitions triplets among processes, which independently sample local mini-batches, fetch embeddings, compute gradients, and apply updates.
- Distributed training: Distributed training partitions graph structures and embeddings across machines, making communication reduction central to retrieving and updating remote embeddings.
- Graph partitioning: METIS partitions co-locate entities with their incident triplets, placing most triplets in diagonal blocks and reducing distributed communication.
- Negative sampling: Joint negative sampling reduces per-batch entity access from O(bd(k + 1)) to O(bd + bkd/д), reaching about b-times less memory access when д grows as large as b.
- Gradient updates: Sparse relation-gradient updates reduce data transferred during multi-processing, multi-GPU, and distributed training.
4 RELATED WORK
Related work includes packages targeting large-scale knowledge graph embedding through Python, distributed, or multi-GPU approaches, but each has important scalability or hardware constraints.
- OpenKE offers many knowledge graph embedding models but cannot scale to very large graphs because it is implemented entirely in Python.
- Pytorch-BigGraph supports distributed training for large graphs but does not support GPU training and uses a block-based scheduling strategy.
- GraphVite focuses on multi-GPU training without distributed support and reduces data movement at the cost of embedding staleness and slower convergence.
5 EXPERIMENTAL METHODOLOGY
The evaluation compares DGL-KE with existing packages across hardware settings and knowledge graph datasets using link prediction and standard ranking metrics.
- DGL-KE is evaluated against GraphVite and Pytorch-BigGraph on multi-GPU, many-core CPU, and distributed-training configurations.
- The evaluation uses FB15k, Freebase, and WN18, with FB15k and WN18 serving as standard knowledge graph embedding benchmarks.
- For FB15k and WN18, corrupted head and tail entities generate filtered negative triplets for validation and testing.
- Performance is measured on link prediction using Hit@k, Mean Rank, and Mean Reciprocal Rank.
- For FB15k and WN18, extensive hyper-parameter searches report configurations achieving the best MRR, whereas Freebase uses one configuration selected from FB15k.
6.1 Multi-GPU training
DGL-KE’s multi-GPU optimizations improve tensor efficiency, reduce CPU–GPU data movement, and overlap computation with communication. Combined, these techniques provide strong training speedups while retaining accuracy.
- Negative sampling: About 4× speedup comes from joint negative sampling on a single GPU, while 8-GPU training reaches about 40× speedup by reducing embedding transfers.Naive sampling swaps many more entity embeddings between CPU and GPU, making communication the bottleneck.
- Asynchronous updates: Roughly 40% speedup on Freebase results from overlapping CPU and GPU computation, although gains are limited on some small-graph models.The benefit depends on CPU and GPU computation times and is larger for graphs with slow random entity-embedding access.
- Relation partitioning: Relation partitioning generally provides over 10% speedup by pinning relation embeddings in GPUs and reducing CPU–GPU data movement.The gain is especially significant for TransR because its relation-specific projection matrices create substantial communication.
- Overall multi-GPU scaling: DGL-KE accelerates multi-GPU training almost linearly and improves Freebase utilization with 16 processes across 8 GPUs.Two processes per GPU overlap computation and CPU–GPU communication across GPU and PCIe resources.
- Overall speed and accuracy: DGL-KE trains many FB15k models within minutes and trains Freebase models in roughly one or two hours, including TransR in about 8 hours on 8 GPUs.These timings follow the combined application of the evaluated optimization techniques.
- Overall speed and accuracy: With a maximum speedup of 11× in single-GPU training, DGL-KE sacrifices little accuracy under the reported fastest configurations.The comparison uses the same total number of epochs for baseline and fastest settings.
6.2 Many-core training
DGL-KE extends its optimization strategy to many-core CPU training and scales well on a 48-core r5dn instance.
- Many-core training: DGL-KE scales well on an r5dn instance with 48 CPU cores for many-core training.The section also reports TransE and DistMult accuracy with 48 CPU cores.
6.3 Distributed training
Distributed training on Freebase compares single-machine training with random and METIS graph partitioning. METIS provides substantial speedup without reducing model accuracy.
- Distributed training: Nearly 3.5× speedup over the single-machine baseline is achieved with METIS partitioning without sacrificing model accuracy.METIS also delivers about 20% speedup over random partitioning because it produces lower overhead.
6.4 Overall performance
DGL-KE is evaluated against GraphVite and PBG across available models and hardware settings. It is consistently faster while targeting comparable accuracy, with efficiency gains linked to faster convergence and less unnecessary computation.
- Evaluation scope: The evaluation compares DGL-KE with GraphVite and PBG only on models available in those competing packages.Experiments cover both CPU and GPU settings using the datasets listed in Table 3.
- Comparison with GraphVite: DGL-KE is consistently faster than GraphVite on FB15k and WN18 when training models to reach similar accuracy.For most models, DGL-KE is 5× faster because it converges in fewer than 100 epochs while GraphVite requires thousands.
- Comparison with PBG: DGL-KE runs twice as fast as PBG when training KGE models on Freebase.PBG treats relation embeddings as dense weights, involving all graph relations in a batch; DGL-KE reduces the relation embeddings involved.
- Comparison with PBG: DGL-KE’s efficiency advantage over PBG is associated with reducing computation and data movement from unnecessary relation-embedding participation in batches.PBG’s batch computation involves all relation embeddings, reported as 10 times more than necessary on Freebase.
7 CONCLUSIONS
DGL-KE trains knowledge graph embeddings at large scale by improving locality, reducing data communication, and harnessing parallel computing. It significantly outperforms state-of-the-art packages across diverse hardware while scaling almost linearly and maintaining high model accuracy.
- DGL-KE trains knowledge graph embeddings at large scale using optimization techniques that improve locality and reduce data communication.
- DGL-KE harnesses parallel computing capacity across many-core CPUs, multi-GPU systems, and clusters of machines.
- DGL-KE significantly outperforms state-of-the-art knowledge graph embedding packages across varied hardware configurations.
- DGL-KE scales almost linearly with machine resources while maintaining very high model accuracy.