Scaling Knowledge Graph Embedding Models for Link Prediction

Nasrullah Sheikh, Xiao Qin, Berthold Reinwald, Chuan Lei

Abstract: Developing scalable solutions for training Graph Neural Net- works (GNNs) for link prediction tasks is challenging due to the inherent data dependencies which entail high com- putational costs and a huge memory footprint. We propose a new method for scaling training of knowledge graph em- bedding models for link prediction to address these chal- lenges. Towards this end, we propose the following algo- rithmic strategies: self-sufficient partitions, constraint-based negative sampling, and edge mini-batch training. The ex- perimental evaluation shows that our scaling solution for GNN-based knowledge graph embedding models achieves a 16x speed up on benchmark datasets while maintaining a comparable model performance to non-distributed methods on standard metrics.