# SparseTransX: Efficient Training of Translation-Based Knowledge Graph Embeddings Using Sparse Matrix Operations

arXiv: https://arxiv.org/abs/2502.16949 · HTML: https://mlsanity.com/paper/2502.16949/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 27 of its 44 references that are in the map.
- First posted on arXiv 24 Feb 2025.

## Builds on (most cited first, 27 of 27)

- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Modeling Relational Data with Graph Convolutional Networks](https://mlsanity.com/paper/1703.06103/) (2017)
- [Embedding Entities and Relations for Learning and Inference in Knowledge Bases](https://mlsanity.com/paper/1412.6575/) (2014)
- [Complex Embeddings for Simple Link Prediction](https://mlsanity.com/paper/1606.06357/) (2016)
- [RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space](https://mlsanity.com/paper/1902.10197/) (2019)
- [Deep Graph Library: A Graph-Centric, Highly-Performant Package for Graph Neural Networks](https://mlsanity.com/paper/1909.01315/) (2019)
- [A Novel Embedding Model for Knowledge Base Completion Based on Convolutional Neural Network](https://mlsanity.com/paper/1712.02121/) (2017)
- [PyTorch-BigGraph: A Large-scale Graph Embedding System](https://mlsanity.com/paper/1903.12287/) (2019)
- [SimKGC: Simple Contrastive Knowledge Graph Completion with Pre-trained Language Models](https://mlsanity.com/paper/2203.02167/) (2022)
- [DGL-KE: Training Knowledge Graph Embeddings at Scale](https://mlsanity.com/paper/2004.08532/) (2020)
- [A Semantic Matching Energy Function for Learning with Multi-relational Data](https://mlsanity.com/paper/1301.3485/) (2013)
- [SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning](https://mlsanity.com/paper/2207.04606/) (2022)
- [TorusE: Knowledge Graph Embedding on a Lie Group](https://mlsanity.com/paper/1711.05435/) (2017)
- [FeatGraph: A Flexible and Efficient Backend for Graph Neural Network Systems](https://mlsanity.com/paper/2008.11359/) (2020)
- [TransG : A Generative Mixture Model for Knowledge Graph Embedding](https://mlsanity.com/paper/1509.05488/) (2015)
- [Differentiating Concepts and Instances for Knowledge Graph Embedding](https://mlsanity.com/paper/1811.04588/) (2018)
- [FusedMM: A Unified SDDMM-SpMM Kernel for Graph Embedding and Graph Neural Networks](https://mlsanity.com/paper/2011.06391/) (2020)
- [TransA: An Adaptive Approach for Knowledge Graph Embedding](https://mlsanity.com/paper/1509.05490/) (2015)
- [Link Prediction with Attention Applied on Multiple Knowledge Graph Embedding Models](https://mlsanity.com/paper/2302.06229/) (2023)
- [Pykg2vec: A Python Library for Knowledge Graph Embedding](https://mlsanity.com/paper/1906.04239/) (2019)
- [TorchKGE: Knowledge Graph Embedding in Python and PyTorch](https://mlsanity.com/paper/2009.02963/) (2020)
- [Force2Vec: Parallel force-directed graph embedding](https://mlsanity.com/paper/2009.10035/) (2020)
- [GraphStorm: all-in-one graph machine learning framework for industry applications](https://mlsanity.com/paper/2406.06022/) (2024)
- [BatchLayout: A Batch-Parallel Force-Directed Graph Layout Algorithm in Shared Memory](https://mlsanity.com/paper/2002.08233/) (2020)
- [iSpLib: A Library for Accelerating Graph Neural Networks using Auto-tuned Sparse Operations](https://mlsanity.com/paper/2403.14853/) (2024)
- [Distributed-Memory Parallel Algorithms for Sparse Matrix and Sparse Tall-and-Skinny Matrix Multiplication](https://mlsanity.com/paper/2408.11988/) (2024)

## Abstract (arXiv)

Knowledge graph (KG) learning offers a powerful framework for generating new knowledge and making inferences. Training KG embedding can take a significantly long time, especially for larger datasets. Our analysis shows that the gradient computation of embedding is one of the dominant functions in the translation-based KG embedding training loop. We address this issue by replacing the core embedding computation with SpMM (Sparse-Dense Matrix Multiplication) kernels. This allows us to unify multiple scatter (and gather) operations as a single operation, reducing training time and memory usage. We create a general framework for training KG models using sparse kernels and implement four models, namely TransE, TransR, TransH, and TorusE. Our sparse implementations exhibit up to 5.3x speedup on the CPU and up to 4.2x speedup on the GPU with a significantly low GPU memory footprint. The speedups are consistent across large and small datasets for a given model. Our proposed sparse approach can be extended to accelerate other translation-based (such as TransC, TransM, etc.) and non-translational (such as DistMult, ComplEx, RotatE, etc.) models as well. An implementation of the SpTransX framework is publicly available as a Python package in https://github.com/HipGraph/SpTransX.