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

要約

ナレッジグラフ（KG）学習は、新しい知識を生成し、推論を行うための強力なフレームワークを提供します。
KG埋め込みのトレーニングには、特により大きなデータセットでは、かなり長い時間がかかる場合があります。
私たちの分析は、埋め込みの勾配計算が翻訳ベースのKG埋め込みトレーニングループの支配的な機能の1つであることを示しています。
コア埋め込み計算をSPMM（スパース密度のマトリックス増殖）カーネルに置き換えることにより、この問題に対処します。
これにより、複数の散布（および収集）操作を単一の操作として統合し、トレーニング時間とメモリの使用量を削減できます。
スパースカーネルを使用してKGモデルをトレーニングするための一般的なフレームワークを作成し、4つのモデル、つまりTranse、TransR、Transh、およびToruseを実装します。
当社のスパースの実装は、CPUで最大5.3倍のスピードアップ、GPUで最大4.2倍のスピードアップを示し、GPUメモリフットプリントが大幅に低くなります。
特定のモデルの大小のデータセットでスピードアップは一貫しています。
提案されたスパースアプローチは、他の翻訳ベース（transc、transmなど）および非翻訳（Distmult、複雑、回転など）モデルを加速するように拡張できます。
sptransxフレームワークの実装は、https：//github.com/hipgraph/sptransxでPythonパッケージとして公開されています。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google