[論文レビュー] SMORE: Knowledge Graph Completion and Multi-hop Reasoning in Massive Knowledge Graphs

Hongyu Ren, Hanjun Dai+5 more|arXiv (Cornell University)|Oct 28, 2021

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

ひとこと要約

SMORE は、大規模知識グラフにおける単一ホップの知識グラフ補完とマルチホップ推論の両方をスケーラブルに実行するフレームワークであり、双方向の破棄サンプリングと非同期 CPU-GPU 計算により、最先端のトレーニングスループットを達成する。Freebase（86Mエンティティ、338Mエッジ）上でマルチホップ推論を実行可能であり、従来のフレームワークと比較して2.2倍の高速化を達成し、8GPUでほぼ線形のスケーリングを実現。GPUメモリはわずか2GBで実現可能である。

ABSTRACT

Knowledge graphs (KGs) capture knowledge in the form of head--relation--tail triples and are a crucial component in many AI systems. There are two important reasoning tasks on KGs: (1) single-hop knowledge graph completion, which involves predicting individual links in the KG; and (2), multi-hop reasoning, where the goal is to predict which KG entities satisfy a given logical query. Embedding-based methods solve both tasks by first computing an embedding for each entity and relation, then using them to form predictions. However, existing scalable KG embedding frameworks only support single-hop knowledge graph completion and cannot be applied to the more challenging multi-hop reasoning task. Here we present Scalable Multi-hOp REasoning (SMORE), the first general framework for both single-hop and multi-hop reasoning in KGs. Using a single machine SMORE can perform multi-hop reasoning in Freebase KG (86M entities, 338M edges), which is 1,500x larger than previously considered KGs. The key to SMORE's runtime performance is a novel bidirectional rejection sampling that achieves a square root reduction of the complexity of online training data generation. Furthermore, SMORE exploits asynchronous scheduling, overlapping CPU-based data sampling, GPU-based embedding computation, and frequent CPU--GPU IO. SMORE increases throughput (i.e., training speed) over prior multi-hop KG frameworks by 2.2x with minimal GPU memory requirements (2GB for training 400-dim embeddings on 86M-node Freebase) and achieves near linear speed-up with the number of GPUs. Moreover, on the simpler single-hop knowledge graph completion task SMORE achieves comparable or even better runtime performance to state-of-the-art frameworks on both single GPU and multi-GPU settings.

研究の動機と目的

• 大規模知識グラフにおけるマルチホップ推論のためのスケーラブルなフレームワークの欠如に対処すること。
• すべてのトレーニングインスタンスを事前に生成しないで、マルチホップクエリのための効率的なオンライントレーニングデータ生成を可能にすること。
• 大規模KGにおける高スループットと低メモリ使用量を実現するため、アルゴリズム的およびシステムレベルの設計を共同最適化すること。
• 統一的でスケーラブルなフレームワーク内で、単一ホップとマルチホップ推論の両方のタスクをサポートすること。

提案手法

• オンライントレーニングデータ生成の複雑さを平方根の要因で低減するため、双方向の破棄サンプリングを導入する。
• CPUベースのデータサンプリングとGPUベースの埋め込み計算、および頻回なCPU-GPU I/Oを重ねて実行する非同期スケジューリングを採用する。
• 関係を跨ぐマルチホップパスの整合性を保つために、単一マシンで非パーティショニングされたグラフ設計を採用する。
• スパースな埋め込みとオプティマイザーの更新を適用することで、トレーニング中のGPUメモリ使用量を最小限に抑える。
• GQE、Q2B、BetaE などの複数のクエリ埋め込みモデルをサポートし、単一GPUおよびマルチGPU環境で統一されたトレーニングパイプラインを提供する。
• 全探索を避けるために双方向の走査を用いて負例エンティティの効率的サンプリングを実現し、スケーラビリティを向上させる。

実験結果

リサーチクエスチョン

• RQ1 8000万エンティティ以上、3億エッジ以上の大規模知識グラフにおいて、マルチホップ推論を効率的にスケーリングすることは可能か？
• RQ2 マルチホップクエリのためのオンライントレーニングデータ生成を、GPUの計算能力を最大限に活用できるほど効率的に実現できるか？
• RQ3 統一されたフレームワークが、最小限のシステムオーバーヘッドで単一ホップの知識グラフ補完と複雑なマルチホップ推論の両方をサポートできるか？
• RQ4 大規模KGで高スループットと低メモリ使用量を達成するために、どのようなシステム最適化が必要か？
• RQ5 双方向サンプリングは、全探索やランダムサンプリングと比較して、トレーニング効率とモデル性能の面で優れているか？

主な発見

• SMORE は、Freebase で400次元の埋め込みをトレーニングする際、GPUメモリをわずか2GBで使用しながらも、従来のマルチホップKGフレームワークと比較して2.2倍の高いスループットを達成した。
• GPU数の増加に伴いトレーニング速度がほぼ線形にスケーリングされ、通信オーバーヘッドが最小限で、非同期設計が効果的に機能していることが示された。
• ogbl-wikikg2 データセットでは、BetaE がベースラインの中で最高のパフォーマンスを示した一方、Q2B は FB400k および Freebase で最先端の結果を達成した。
• 双方向サンプリングは、全探索に匹敵またはそれを上回る性能を発揮したが、大規模グラフではトレーニング時間を月単位から数日単位に短縮した。
• 非パーティショニングされたグラフ設計であるにもかかわらず、Marius や PBG などの最先端のフレームワークと同等またはそれ以上の単一ホップリンク予測パフォーマンスを達成した。
• システムは Freebase（86Mノード、338Mエッジ）にまでスケーリング可能であり、これはマルチホップ推論でこれまでにサポートされたKGよりも1500倍も大きい規模である。