用稀疏矩阵运算加速知识图谱嵌入训练，速度提升超5倍。
SparseTransX: Efficient Training of Translation-Based Knowledge Graph Embeddings Using Sparse Matrix Operations
- 用稀疏矩阵乘法替代传统嵌入计算，统一散列与聚集操作。
- 在CPU上最高提速5.3倍，GPU上最高提速4.2倍，显存占用极低。
- 适用于多种翻译型及非翻译型模型，开源可用。
知识图谱学习为生成新知识和推理提供了强大框架，但嵌入训练耗时较长，尤其在大规模数据集上。分析表明，嵌入的梯度计算是翻译型知识图谱嵌入训练中的主要开销。本文提出使用稀疏-稠密矩阵乘法（SpMM）内核替代核心嵌入计算，将多个散列（scatter）和聚集（gather）操作合并为单次操作，显著降低训练时间和内存占用。我们构建了一个通用框架，实现TransE、TransR、TransH和TorusE四种模型。稀疏版本在CPU上最高提速5.3倍，GPU上最高提速4.2倍，且显存消耗极低。该方法在大小数据集上均表现稳定，并可扩展至其他翻译型（如TransC、TransM）及非翻译型模型（如DistMult、ComplEx、RotatE）。代码已开源，可在https://github.com/HipGraph/SpTransX获取。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。