[論文レビュー] Sketch-GNN: Scalable Graph Neural Networks with Sublinear Training Complexity

Mucong Ding, Tahseen Rabbani+3 more|arXiv (Cornell University)|Jun 21, 2024

Advanced Graph Neural Networks被引用数 7

ひとこと要約

Sketch-GNNは隣接行列とノード特徴のコンパクトなスケッチ上でGNNを訓練し、グラフサイズに対してサブリニアの訓練時間とメモリを達成しつつ、競争力のある精度を維持します。多項式テンソルスケッチと学習可能なLSHを用いて、非線形活性化とメッセージパッシングを近似します。

ABSTRACT

Graph Neural Networks (GNNs) are widely applied to graph learning problems such as node classification. When scaling up the underlying graphs of GNNs to a larger size, we are forced to either train on the complete graph and keep the full graph adjacency and node embeddings in memory (which is often infeasible) or mini-batch sample the graph (which results in exponentially growing computational complexities with respect to the number of GNN layers). Various sampling-based and historical-embedding-based methods are proposed to avoid this exponential growth of complexities. However, none of these solutions eliminates the linear dependence on graph size. This paper proposes a sketch-based algorithm whose training time and memory grow sublinearly with respect to graph size by training GNNs atop a few compact sketches of graph adjacency and node embeddings. Based on polynomial tensor-sketch (PTS) theory, our framework provides a novel protocol for sketching non-linear activations and graph convolution matrices in GNNs, as opposed to existing methods that sketch linear weights or gradients in neural networks. In addition, we develop a locality-sensitive hashing (LSH) technique that can be trained to improve the quality of sketches. Experiments on large-graph benchmarks demonstrate the scalability and competitive performance of our Sketch-GNNs versus their full-size GNN counterparts.

研究の動機と目的

• 全グラフのメモリを必要とせず、グラフサイズに対する線形時間依存性を持たない非常に大規模なグラフ上でのGNN訓練を動機づける。
• 訓練前に隣接行列と特徴行列を圧縮するスケッチベースのフレームワークを紹介する。
• 非線形活性化をGNNに適用するために polynomial tensor sketching を拡張する。
• 訓練中にスケッチを適応させる学習可能な locality-sensitive hashing を開発して精度を向上させる。

提案手法

• CXW 上の非線形演算としてメッセージパッシングを表現し、CXW を polynomial tensor sketches で近似する。
• count sketchとtensor sketchを用いてCと XWを小さな行列にスケッチし、サブリニアなメモリと時間を実現する。
• 非線性をべき級数で展開し、係数 ck を学習して活性化関数を近似する。
• 訓練を通じてスケッチを適応させる学習可能な SimHash ベースの LSH を導入する。
• 層間でスケッチ表現を効率的に計算できる2面スケッチフレームワークを提供する。
• 層間の誤差蓄積を軽減するためにスキップ接続を組み込む。

実験結果

リサーチクエスチョン

• RQ1 隣接行列と特徴行列をスケッチして、グラフサイズに対してサブリニアな時間とメモリでGNN訓練を実現できるか？
• RQ2 完全空間に戻さず、スケッチベースのフレームワーク内でGNNの非線形活性化をどのように近似できるか？
• RQ3 訓練中に学習可能なLSHがスケッチの品質を改善し、競争力のある精度を維持できるか？
• RQ4 Sketch-GNNのメモリ・時間・精度の理論的・経験的トレードオフは、一般的なGNNアーキテクチャ間でどうなるか？

主な発見

• Sketch-GNNはグラフサイズに対してサブリニアな訓練計算量を達成し、巨大なベンチマークでも競争力のある精度を維持する。
• 本フレームワークは n×n の隣接行列と n×d の特徴行列をはるかに小さなスケッチにスケッチし、層ごとに O(c) のメモリと時間を可能にする（c は n に依存しない）。
• Polynomial tensor sketch (PTS) を用いて、CXと XW の tensor/count sketches と組み合わせて非線形活性化を近似する。
• 学習可能な LSH アプローチは訓練中にハッシュテーブルを適応させ、スケッチ品質を改善しパフォーマンス低下を低減する。
• スキップ接続は、スケッチベースのパイプライン全体での層間の誤差蓄積を緩和するのに役立つ。
• 実験結果は、フルグラフGNNと比較して競争力のある性能と、大規模グラフでの有利なスケーラビリティを示している。