SparseTransX：使用稀疏矩阵运算高效训练基于翻译的知识图谱嵌入
机器学习
2025-05-01 v3 计算与语言
摘要
知识图谱（KG）学习为生成新知识和进行推理提供了一个强大的框架。训练 KG 嵌入可能需要非常长的时间，特别是对于较大的数据集。我们的分析表明，嵌入的梯度计算是基于翻译的 KG 嵌入训练循环中的主要函数之一。我们通过用 SpMM（稀疏-稠密矩阵乘法）内核替换核心嵌入计算来解决这个问题。这使我们能够将多个 scatter 和 gather 操作统一为单一操作，从而减少训练时间和内存使用。我们创建了一个使用稀疏内核训练 KG 模型的通用框架，并实现了四个模型，即 TransE、TransR、TransH 和 TorusE。我们的稀疏实现在 CPU 上实现了高达 5.3 倍的加速，在 GPU 上实现了高达 4.2 倍的加速，且 GPU 内存占用极低。对于给定的模型，这种加速在不同大小数据集上保持一致。我们提出的稀疏方法也可以扩展到加速其他基于翻译的模型（如 TransC、TransM 等）和非翻译模型（如 DistMult、ComplEx、RotatE 等）。SpTransX 框架的实现已作为 Python 包公开于 https://github.com/HipGraph/SpTransX。
引用
@article{arxiv.2502.16949,
title = {SparseTransX: Efficient Training of Translation-Based Knowledge Graph Embeddings Using Sparse Matrix Operations},
author = {Md Saidul Hoque Anik and Ariful Azad},
journal= {arXiv preprint arXiv:2502.16949},
year = {2025}
}
备注
15 pages, 9 figures, 9 tables. MLSys 2025