Scaling knowledge graph embedding models for link prediction

Nasrullah Sheikh, Xiao Qin, Berthold Reinwald, Chuan Lei

Abstract: Developing scalable solutions for training Graph Neural Networks (GNNs) for link prediction tasks is challenging due to the inherent data dependencies which entail high computational costs and a huge memory footprint. We propose a new method for scaling training of knowledge graph embedding models for link prediction to address these challenges. Towards this end, we propose the following algorithmic strategies: self-sufficient partitions, constraint-based negative sampling, and edge mini-batch training. The experimental 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.