[論文レビュー] 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 を提供する。TorchKGE はリンク予測の評価を、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 埋め込みのコアタスクであるリンク予測のための高パフォーマンスな評価パイプラインを提供し、学習ワークフローにおけるボトル neck を軽減する。
• 最小限の依存関係、きめ細かいドキュメンテーション、コミュニティサポートを備えた洗練された PyTorch ネイティブ API を提供し、研究者やエンジニアの参入障壁を低減する。
• モジュラー設計とオール・オブ・ザ・ボックス実装を活用して、新しい KG 埋め込みアーキテクチャの迅速なプロトタイピングとベンチマークを可能にする。

提案手法

• すべてのモデルを PyTorch の nn.Module のサブクラスとして実装し、自動微分、GPU 加速、最適化ツールを活用する。
• リンク予測のための専用で高度に最適化された評価モジュールを設計し、バッチ処理およびベクトル化された方法で、すべての候補エンティティのスコアを計算する。
• 学習、検証、テスト分割を管理するためのモジュラーなデータ構造（KnowledgeGraph クラス）を提供し、設定可能な分割戦略をサポートする。
• 一元化されたインターフェースを通じて、複数のネガティブサンプリング戦略（一様、ベルヌーイ、位置ベース）を統合し、ネガティブハード例を効率的に使用した学習を可能にする。
• 広範なアクセス性とコミュニティ貢献を確保するため、PyPI を用いた配布と BSD ライセンスを採用する。
• utils.datasets モジュール内のユーティリティ関数を通じて、標準ベンチマークデータセット（例：FB15k, WN18, YAGO3-10）の読み込みをサポートする。

実験結果

リサーチクエスチョン

• RQ1 異なる研究グループ間での公平なモデル比較を可能にするために、統一的で効率的かつ再現可能な知識グラフ埋め込みフレームワークをどのように設計できるか？
• RQ2 PyTorch ネイティブ実装においてスコアリングおよびランク付けパイプラインを最適化することで、リンク予測の評価でどの程度のパフォーマンス向上が達成できるか？
• RQ3 高度に最適化された評価モジュールは、学習のオーバーヘッドをどの程度低減し、ハイパーパramータチューニングの効率を向上させるか？
• RQ4 最小限の依存関係を備えた PyTorch ベースのライブラリは、トレーニング速度において競争力を持ちつつ、既存のフレームワークに比べて評価遅延で顕著な優位性を示せるか？
• RQ5 TorchKGE におけるリンク予測評価のパフォーマンスは、同じハードウェアおよびハイパーパramータ設定下で、OpenKE や AmpliGraph と比較してどの程度か？

主な発見

• WN18 で RESCAL を評価した際、TorchKGE は OpenKE より 24 倍高速で、評価時間は 7.06 秒（TorchKGE）対 178.4 秒（OpenKE）であった。
• テストされたすべてのモデルとデータセットにおいて、TorchKGE は平均して OpenKE や AmpliGraph より少なくとも 3 倍速いリンク予測評価を達成した。
• TorchKGE のトレーニングエポック時間は OpenKE や AmpliGraph と比較して 2 倍以内であり、評価速度に注力しているにもかかわらず、強力なトレーニング効率を示している。
• TorchKGE の評価モジュールは非常に最適化されており、FB15k で ComplEx を使用したリンク予測の評価時間は、TorchKGE では 96.4 秒、AmpliGraph では 354.7 秒にまで延びた。
• TorchKGE は、主な KG 埋め込みモデル（TransE, TransD, RESCAL, ComplEx）をすべてサポートし、オリジナルの論文への参照付きで一貫性があり、文書化された実装を提供する。
• このライブラリのパフォーマンス優位性は、評価ワークロードにおいて顕著であり、TorchKGE は OpenKE や AmpliGraph より最大 24 倍速く、トレーニング中の頻繁な評価が現実的かつ効率的に行えるようになる。