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

Research output: Chapter in Book/Report/Conference proceedingConference contribution

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 2× ∼ 5× speedup over the best competing approaches. DGL-KE is available on https://github.com/awslabs/dgl-ke.

Original languageEnglish (US)
Title of host publicationSIGIR 2020 - Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval
PublisherAssociation for Computing Machinery, Inc
Pages739-748
Number of pages10
ISBN (Electronic)9781450380164
DOIs
StatePublished - Jul 25 2020
Externally publishedYes
Event43rd Annual International ACM SIGIR Conference on Research and Development in Information Retrieval, SIGIR 2020 - Virtual, Online, China
Duration: Jul 25 2020Jul 30 2020

Publication series

NameSIGIR 2020 - Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval

Conference

Conference43rd Annual International ACM SIGIR Conference on Research and Development in Information Retrieval, SIGIR 2020
CountryChina
CityVirtual, Online
Period7/25/207/30/20

Keywords

  • distributed training
  • knowledge graph embeddings
  • large scale

Fingerprint Dive into the research topics of 'DGL-KE: Training Knowledge Graph Embeddings at Scale'. Together they form a unique fingerprint.

Cite this