Scaling Weisfeiler-Leman Expressiveness Analysis to Massive Graphs with GPUs

技術要約：GPUを用いた大規模グラフへのWeisfeiler–Leman表現力解析のスケーリング

問題定義
Weisfeiler–Leman (1-WL) テストとその関連する安定彩色（stable coloring）は、グラフニューラルネットワーク（GNN）の表現力を理解するための基礎となる。1-WL安定彩色は、グラフ構造を区別する能力の上限を提供し、同じ安定色を受け取ったノードは、そのようなネットワークにとって区別不能となる。しかし、実世界の巨大なグラフに対してこの安定彩色を計算することは、以下の2つの決定的なボトルネックに直面する：

1. 逐次実行: 分割細分化（partition refinement）に基づく古典的なアルゴリズム（例：$O(m \log n)$ のスプリッターベースの手法）は本質的に逐次的であり、現代のGPUアーキテクチャが持つ大規模な並列性を活用できない。
2. グローバルメモリ要件: これらのアルゴリズムは、グラフ全体が同時にメモリ内に存在することを要求する。この「グローバル」な性質により、単一のデバイスや大規模なCPUクラスターの容量を超えるウェブスケールのグラフ（例：数十億のエッジを持つもの）の解析が困難となっている。

並列化や線形化を試みた既存の試みの多くは、浮動小数点演算に依存しており、大規模なスケールでは数値的不安定性を引き起こすか、あるいはデバイスメモリを超えるグラフのメモリ制約に対処できていない。

手法
著者らは、1-WL安定彩色の線形代数的な解釈を利用して、計算のスケーラビリティとメモリのスケーラビリティを切り離す二段構えのアプローチを提案している。

1. ランダム化された線形代数的細分化（計算のスケーラビリティ）
逐次的な分割細分化の代わりに、著者らは等価分割（equitable partition）の不変性を利用して問題を再定式化する。すなわち、ある分割は、そのブロック一定（block-constant）なベクトルが隣接行列の下でブロック一定であり続ける場合にのみ安定である。

• アルゴリズム: この手法は、現在の分割に対して対称なランダムベクトルをサンプリングし、それらを隣接行列と乗算し（$z = Aw$）、結果として得られるベクトルの対称性に基づいて分割を精緻化するプロセスを反復する。
• 算術ドメイン: 浮動小数点を用いたべき乗反復（power iteration）が引き起こす数値的不安定性（大規模なグラフにおいて誤ったマージを引き起こす可能性がある）を避けるため、アルゴリズムはすべて符号なし64ビット整数演算を用いた剰余環 $\mathbb{Z}/2^{64}\mathbb{Z}$ 内で動作する。
• 理論的保証: 隣接行列の行和が $2^{64}/2$ によって抑えられるという構造的制約を利用することで、著者らは確率的な正当性の保証を導出している。ランダムなプローブは、最大 $10^{-9}$ の失敗確率で、欠落した精緻化を検出できる。これは、体（field）を必要とする古典的なシュワルツ・ジッペル補題を回避するものである。
• GPUマッピング: コアとなる操作は、単一の疎行列ベクトル積（SpMV）と辞書式ソートに還元され、これらはGPUハードウェア上で非常に効率的に動作する。

2. 正当性を保持するバッチ処理（メモリのスケーラビリティ）
デバイスメモリに収まらないグラフを扱うために、著者らは、グローバルな安定彩色の正当性を維持しながら、エッジ集合を互いに素な部分集合に分解するバッチングスキームを導入している。

• 内部ノードと境界ノード: 与えられたエッジのバッチに対して、ノードは「内部ノード」（すべての流出エッジがバッチ内に存在する）または「境界ノード」（流出エッジがバッチの境界を越える）として分類される。
• 局所的な精緻化: アルゴリズムは各バッチを独立して処理する。重要な点として、各バッチの初期分割では、すべての内部ノードを単一のブロックに配置し、すべての境界ノードをシングルトン（単独）のブロックに配置する。これにより、精緻化は、そのバッチ内で完全な近傍情報が利用可能なノードのみをマージすることを保証する。
• 反復的な縮約: 局所的な精緻化はグローバルな分割へと統合される。その後、アルゴリズムは商グラフ（現在の分割のブロックをノードとするグラフ）を構築し、プロセスを繰り返す。
• 停止条件: グラフが単一のバッチに収まるまで（最も粗い安定彩色が得られる）、あるいは不動点に達するまで、このプロセスは反復される。もし不動点に達する前にグラフがメモリに収まった場合、その結果は最も粗いものに対する精緻化としての有効な安定彩色となる。

主な貢献

1. ランダム化された精緻化アルゴリズム: 整数演算で動作し、タイトな確率的保証（誤差 $<10^{-9}$）を持つ、1-WL安定彩色のためのモンテカルロ・フレームワーク。これは、浮動小数点法による精度損失を回避する。
2. バッチングスキーム: 独立したエッジ部分集合を処理し、商グラフを介して結果をマージすることで、デバイスメモリを超えるグラフ上での1-WL計算を可能にする分解戦略。
3. GPU実装: これらのアルゴリズムを SpMV、ソート、リダクションといった GPU プリミティブに直接マッピングする CUDA ベースの実装。これにより、CPU ベースラインに対して大幅な高速化を実現した。

実験結果
著者らは、11個の中規模グラフと3個のウェブスケールグラフ（330億〜480億エッジ）を用いて、彼らの手法を評価した。

• 正当性: メモリに全グラフが収まる中規模グラフ（$N_b=1$）において、ランダム化アルゴリズムは決定論的なベースラインと一致する、正確な最も粗い安定彩色を復元した。
• 数値的安定性: 提案された整数演算による手法は、大規模なグラフにおいて（以前の浮動小数点べき乗反復法によって）誤ってマージされていたノードを正常に区別した（例：it-2004 データセットでは、80万個以上のブロックの相違が見られた）。
• バッチングの品質: メモリ制約をシミュレートするために複数のバッチ（$N_b=2, 4$）を使用した際、得られた分割は経験的に最も粗い安定彩色の5%以内に収まった。いくつかのデータセットでは、その差はゼロであった。
• パフォーマンス:
  • 速度: シングルGPU 実装は、古典的なシングルスレッド CPU 分割細分化に対して最大 138倍、マルチコア CPU ベースラインに対して最大 449倍 の高速化を達成した。
  • スケーラビリティ: 300億エッジを超えるウェブスケールグラフにおいて、すべての CPU ベースライン（マルチコアの記号的縮約を含む）は、タイムアウトするかメモリ不足となった。提案手法はこれらのインスタンスの計算に成功し、ブロック数の約39〜46%の削減を達成した。

意義と主張
本論文は、300億エッジを超えるウェブスケールグラフ上で1-WL安定彩色を計算することに成功した最初の事例であると主張している。これは、メモリと計算の制限により以前は不可能であったタスクである。その意義は以下にある：

• 表現力解析の実現: 小規模または中規模のグラフに限定されることなく、実世界の巨大なデータセットにおける GNN の表現力の経験的解析を可能にする。
• アーキテクチャへの適合: 本来は逐次的かつグローバルなものと見なされていた 1-WL 問題が、GPU アーキテクチャに合わせて調整された線形代数的定式化とバッチング戦略を用いることで、効果的に並列化および分散化できることを示している。
• 堅牢性: 整数演算による数値的安定性と、バッチングによるメモリのスケーラビリティの両方を備えた方法を提供し、従来の近似手法の失敗モードを克服している。

著者らは、本手法が以前は到達不可能であったスケールでの 1-WL 表現力解析を可能にすると結論付けており、安定彩色と等価分割の対応関係が、ダウンストリームのグラフ ML タスクや非線形力学系のモデル縮約への応用における将来的な道筋を提供することを指摘している。