[論文レビュー] DGL-KE: Training Knowledge Graph Embeddings at Scale

Da Zheng, Xiang Song+7 more|arXiv (Cornell University)|Apr 18, 2020

Advanced Graph Neural Networks参考文献 17被引用数 36

ひとこと要約

DGL-KE は大規模グラフ上の知識グラフ埋め込みのスケーラブルな訓練のための新しい最適化を導入し、GraphVite と Pytorch-BigGraph に対してマルチ-GPU、Many-core CPU、分散設定で大幅な高速化を達成します。

ABSTRACT

Knowledge graphs have emerged as a key abstraction for organizing information in diverse domains and their embeddings are increasingly used to harness their information in various information retrieval and machine learning tasks. However, the ever growing size of knowledge graphs requires computationally efficient algorithms capable of scaling to graphs with millions of nodes and billions of edges. This paper presents DGL-KE, an open-source package to efficiently compute knowledge graph embeddings. DGL-KE introduces various novel optimizations that accelerate training on knowledge graphs with millions of nodes and billions of edges using multi-processing, multi-GPU, and distributed parallelism. These optimizations are designed to increase data locality, reduce communication overhead, overlap computations with memory accesses, and achieve high operation efficiency. Experiments on knowledge graphs consisting of over 86M nodes and 338M edges show that DGL-KE can compute embeddings in 100 minutes on an EC2 instance with 8 GPUs and 30 minutes on an EC2 cluster with 4 machines with 48 cores/machine. These results represent a 2x~5x speedup over the best competing approaches. DGL-KE is available on https://github.com/awslabs/dgl-ke.

研究の動機と目的

• 巨大な知識グラフで何百万ものノードと何十億のエッジを持つスケーラブルな埋め込み訓練を動機付ける。
• データの局所性を改善し、通信を削減し、KGE訓練の計算を重ね合わせる最適化を開発・実装する。
• DGL, PyTorch/MXNet, および分散KVStore上に構築されたオープンソースパッケージ（DGL-KE）を提供し、大規模訓練を可能にする。

提案手法

• 知識グラフを計算ユニット間で分割して、跨パーティションのデータ移動を最小化する。
• METISベースのグラフ分割を使用して埋め込みをトリップレットと同じ場所に配置し、機間通信を削減する。
• ジョイントおよび次数ベースのネガティブサンプリングを実装してデータ移動を削減し、より難しいネガティブを生成する。
• 関係埋め込みを分割し、疎な読み取り/更新を伝搬させて多くの関係のデータ転送を最小化する。
• 勾配更新とバッチ処理を重ね合わせてCPU-GPUの遅延を隠し、スループットを向上させる。
• 定期的な同期と分散キーバリューストアを導入してマシン間の訓練を調整する。

実験結果

リサーチクエスチョン

• RQ1 ノード百万〜億規模のグラフで、知識グラフ埋め込みをスケールさせて効率的に訓練するにはどうすればよいか？
• RQ2 速度と埋め込み品質の最良のトレードオフを得るために、グラフ分割、ネガティブサンプリング、データ配置のどの組み合わせが最適か？
• RQ3 非同期、マルチプロセッシング、および分散訓練は、データ転送とメモリのボトルネックを減らしつつ競争力のある埋め込み品質を達成できるか？
• RQ4 DGL-KE の性能は大規模知識グラフに対して既存ツール（GraphVite、PBG）と比べてどうか？
• RQ5 関係分割と勾配更新の重ね合わせが訓練スループットとモデル精度に与える実用的影響は？

主な発見

• DGL-KE は顕著な高速化を達成: 例えば、86Mノード・338Mエッジのグラフで8-GPU EC2インスタンスでは100分、4マシンクラスタ（各マシン48コア）では30分、競合手法と比較して。
• ジョイントネガティブサンプリングはデータ移動を最大で約40倍削減し、テンソル演算の効率を向上させる。
• 次数ベースのネガティブサンプリングはFreebaseのような大規模グラフでモデル精度を向上させるが、訓練速度の向上にはつながらない場合がある。
• CPUとGPU間の勾配更新の重ね合わせは、多くの大規模グラフモデルで約40%の速度アップを実現し、精度を維持する。
• 関係分割はデータ移動を大幅に削減し、特に大きな関係関連のパラメータを持つモデル（例: TransR）で顕著な高速化をもたらす。
• 総じて、DGL-KE はGPU数にほぼ線形にスケールし、大規模グラフ上で複雑なKGEモデルを数時間で訓練でき、GraphVite および PBG を報告された実験で上回る。