Fast and Accurate Learning of Knowledge Graph Embeddings at Scale

Knowledge Graph Embedding (KGE) is used to represent the entities and relations of a KG in a low dimensional vector space. KGE can then be used in a downstream task such as entity classification, link prediction and knowledge base completion. Training on large KG datasets takes a considerable amount of time. This work proposes three strategies which lead to faster training in distributed setting. The first strategy is a reduced communication approach which decreases the All-Gather size by sparsifying the Sparse Gradient Matrix (SGM). The second strategy is a variable margin approach that takes advantage of reduced communication for lower margins but retains the accuracy as obtained by the best fixed margin. The third strategy is called DistAdam which is a distributed version of the popular Adam optimization algorithm. Combining the three strategies results in reduction of training time for the FB250K dataset from twenty-seven hours on one processing node to under one hour on thirty-two nodes with each node consisting of twenty-four cores.