[论文解读] TorchKGE: Knowledge Graph Embedding in Python and PyTorch

Armand Boschin|arXiv (Cornell University)|Sep 7, 2020

Advanced Graph Neural Networks参考文献 8被引用 7

一句话总结

TorchKGE 是一个基于 PyTorch 的开源 Python 库，用于知识图谱嵌入，为实现和评估 KG 嵌入模型提供了快速、高效且用户友好的 API。它将链接预测评估速度显著提升，最高可达 OpenKE 和 AmpliGraph 的 24 倍，同时支持最先进模型、负采样以及标准化评估指标。

ABSTRACT

TorchKGE is a Python module for knowledge graph (KG) embedding relying solely on PyTorch. This package provides researchers and engineers with a clean and efficient API to design and test new models. It features a KG data structure, simple model interfaces and modules for negative sampling and model evaluation. Its main strength is a very fast evaluation module for the link prediction task, a central application of KG embedding. Various KG embedding models are also already implemented. Special attention has been paid to code efficiency and simplicity, documentation and API consistency. It is distributed using PyPI under BSD license. Source code and pointers to documentation and deployment can be found at https://github.com/torchkge-team/torchkge.

研究动机与目标

• 解决在不同编程语言和深度学习框架之间缺乏统一、高效的 KG 嵌入模型训练与评估框架的问题。
• 通过在相同软硬件条件下一致比较模型，提升 KG 嵌入研究的可复现性。
• 提供一个高性能的链接预测评估流水线，这是 KG 嵋入中的核心任务，通常在训练工作流中成为瓶颈。
• 提供简洁、原生 PyTorch 的 API，依赖项极少，文档完善，并有社区支持，以降低研究人员和工程师的使用门槛。
• 通过模块化设计和开箱即用的实现，支持新 KG 嵌入架构的快速原型设计与基准测试。

提出的方法

• 将所有模型实现为 PyTorch 的 nn.Module 子类，以利用自动微分、GPU 加速和优化工具。
• 设计一个专用的、高度优化的链接预测评估模块，以批处理、向量化的方式计算所有候选实体的得分。
• 提供一个模块化的数据结构（KnowledgeGraph 类），用于管理训练、验证和测试划分，并支持可配置的划分策略。
• 通过统一接口集成多种负采样策略（均匀采样、伯努利采样、位置采样），以支持高效训练包含困难负样本的模型。
• 使用 PyPI 进行分发，并采用 BSD 许可证，以确保广泛的可访问性与社区贡献。
• 通过 utils.datasets 模块中的工具函数，支持标准基准数据集（如 FB15k、WN18、YAGO3-10）的加载。

实验结果

研究问题

• RQ1 如何设计一个统一、高效且可复现的知识图谱嵌入框架，以实现在不同研究团队之间的公平模型比较？
• RQ2 通过在原生 PyTorch 实现中优化评分与排序流水线，链接预测评估可实现多大的性能提升？
• RQ3 一个高度优化的评估模块在多大程度上能减少训练开销并提升超参数调优效率？
• RQ4 一个依赖项极少的基于 PyTorch 的库是否能在保持训练速度竞争力的同时，在评估延迟方面显著优于现有框架？
• RQ5 在相同硬件和超参数设置下，TorchKGE 的链接预测评估性能与 OpenKE 和 AmpliGraph 等成熟库相比如何？

主要发现

• 在 WN18 上评估 RESCAL 时，TorchKGE 相较于 OpenKE 实现了 24 倍的评估速度提升，评估时间分别为 7.06 秒与 178.4 秒。
• 在所有测试的模型和数据集上，TorchKGE 的链接预测评估速度平均至少是 OpenKE 和 AmpliGraph 的三倍。
• TorchKGE 的训练 epoch 时间与 OpenKE 和 AmpliGraph 相当，差距在两倍以内，表明尽管专注于评估速度，其训练效率依然出色。
• TorchKGE 的评估模块高度优化，将 FB15k 数据集上 ComplEx 模型的链接预测时间从 96.4 秒（TorchKGE）降低至 354.7 秒（AmpliGraph）。
• TorchKGE 支持所有主流 KG 嵌入模型（TransE、TransD、RESCAL、ComplEx），并提供一致、有文档记录的实现，附有原始论文参考。
• 该库的性能优势在评估工作负载中最为显著，其评估速度相比 OpenKE 和 AmpliGraph 最高可提升 24 倍，使得在训练过程中频繁评估变得切实可行且高效。