[論文レビュー] Marius: Learning Massive Graph Embeddings on a Single Machine

Jason Mohoney, Roger Waleffe+3 more|arXiv (Cornell University)|Jan 20, 2021

Advanced Graph Neural Networks参考文献 33被引用数 10

ひとこと要約

Mariusは、パーティションキャッシュとバッファを考慮したデータ順序付け（BETA）によるデータ移動最適化を通じて、1台のマシン上で大規模グラフ埋め込みを学習するフレームワークを提案する。これにより、最先端のシステムよりも最大10倍高速に学習が可能となり、GPUメモリ16 GB、CPUメモリ64 GBでのみ、10億エッジを超えるグラフと550 GBのパラメータをスケーリング可能である。

ABSTRACT

We propose a new framework for computing the embeddings of large-scale graphs on a single machine. A graph embedding is a fixed length vector representation for each node (and/or edge-type) in a graph and has emerged as the de-facto approach to apply modern machine learning on graphs. We identify that current systems for learning the embeddings of large-scale graphs are bottlenecked by data movement, which results in poor resource utilization and inefficient training. These limitations require state-of-the-art systems to distribute training across multiple machines. We propose Marius, a system for efficient training of graph embeddings that leverages partition caching and buffer-aware data orderings to minimize disk access and interleaves data movement with computation to maximize utilization. We compare Marius against two state-of-the-art industrial systems on a diverse array of benchmarks. We demonstrate that Marius achieves the same level of accuracy but is up to one order of magnitude faster. We also show that Marius can scale training to datasets an order of magnitude beyond a single machine's GPU and CPU memory capacity, enabling training of configurations with more than a billion edges and 550 GB of total parameters on a single machine with 16 GB of GPU memory and 64 GB of CPU memory. Marius is open-sourced at www.marius-project.org.

研究の動機と目的

• 1台のマシン上で大規模グラフ埋め込み学習におけるデータ移動のボトル neck を解消すること。
• GPUメモリ外での学習における不規則なメモリアクセスパターンによって引き起こされる低効率なGPU利用率を克服すること。
• GPUとCPUのメモリ容量を超えるグラフ埋め込みの学習を、I/Oと計算の重ね合わせ最適化によって可能にすること。
• 既存の産業用システムと比較して、大幅に高速化されたトレーニングスループットを維持しながら、高い精度を保つシステム設計をすること。
• I/Oを最小限に抑え、ハードウェア利用率を最大化するための新規なデータ順序付け方式（BETA）とパイプラインアーキテクチャを設計すること。

提案手法

• I/O遅延を隠すために、データ移動、転送、計算をインタリーブするパイプラインアーキテクチャを導入する。
• 頻繁にアクセスされるノード埋め込みパーティションを再利用することで、不要なディスク読み込みを減らすパーティションキャッシュを実装する。
• バッファを考慮したデータ順序付け方式BETAを提案し、近くのノードパーティションを持つエッジをグループ化することでI/Oを最小限に抑える。
• ノード埋め込みにはしきい値付きの非同期学習、エッジ埋め込みには同期学習を採用し、異種エッジタイプのサポートを実現する。
• 計算オーバーヘッドを低減するため、効率的なドット積演算を用いたミニバッチ確率的勾配降下法（SGD）を活用する。
• スループットと埋め込み品質のバランスを取るために、バッファサイズ、パーティション数、およびしきい値境界のシステム設定最適化を実施する。

実験結果

リサーチクエスチョン

• RQ1 分散システムを用いずに1台のマシン上で10億エッジを超えるグラフ埋め込み学習をスケーリング可能か？
• RQ2 GPUメモリ外での学習において、I/Oオーバーヘッドを最小限に抑え、高いGPU利用率を達成するにはどうすればよいか？
• RQ3 大規模グラフ埋め込みの学習効率を維持しながらI/Oコストを低減するデータ順序付け戦略は何か？
• RQ4 パーティションキャッシュやパイプライン化といったシステムレベル最適化は、モデル精度を損なわず、どの程度トレーニングスループットを向上できるか？
• RQ5 バッチサイズ、バッファサイズ、しきい値境界といったハイパーパrameterは、トレーニングパフォーマンスと埋め込み品質にどのように影響するか？

主な発見

• Mariusは平均で70%のGPU利用率を達成したのに対し、DGL-KEは10%、PBGは30%未満であり、ハードウェア利用率が顕著に向上した。
• Freebase86mデータセットにおいて、MariusはDGL-KEおよびPBGと同等のモデル精度を維持しながら最大10倍高速に学習を実行した。
• Mariusは、GPUメモリ16 GB、CPUメモリ64 GBでのみ、10億エッジを超えるグラフと最大550 GBの合計モデルパラメータをスケーリング可能である。
• BETA順序付けは、特にパーティション数が多い状況で、ノード埋め込みのディスクアクセスを最小限に抑え、I/Oコストを削減した。
• 大規模グラフでは、バッチサイズを約10,000にすることでスループットが向上し、精度に損なわれることなく効果を発揮したが、しきい値を超えるとその恩恵は薄れた。
• バッファ容量を2乗のスケールで増加させるとディスクスワップ回数が著しく減少し、バッファサイズの最適化がパフォーマンスに極めて重要であることがわかった。