[論文レビュー] Probabilistically Rewired Message-Passing Neural Networks

Chendi Qian, Andrei Manolache+5 more|arXiv (Cornell University)|Oct 3, 2023

Advanced Graph Neural Networks被引用数 4

ひとこと要約

本稿では、正確で微分可能なkサブセットサンプリングを用いて、タスク関連の事前知識に基づき、グラフ内のエッジの追加・削除を学習する微分可能で原理的であるProbabilistically Rewired Message-Passing Neural Networks (PR-MPNNs)を提案する。この手法は表現力の向上、過剰圧縮の緩和および情報伝達の不足の是正を実現し、分子および引用ネットワークベンチマークにおいて、標準的なMPNNやグラフトランスフォーマーを凌駕する、最先端または競争力のある性能を達成する。

ABSTRACT

Message-passing graph neural networks (MPNNs) emerged as powerful tools for processing graph-structured input. However, they operate on a fixed input graph structure, ignoring potential noise and missing information. Furthermore, their local aggregation mechanism can lead to problems such as over-squashing and limited expressive power in capturing relevant graph structures. Existing solutions to these challenges have primarily relied on heuristic methods, often disregarding the underlying data distribution. Hence, devising principled approaches for learning to infer graph structures relevant to the given prediction task remains an open challenge. In this work, leveraging recent progress in exact and differentiable $k$-subset sampling, we devise probabilistically rewired MPNNs (PR-MPNNs), which learn to add relevant edges while omitting less beneficial ones. For the first time, our theoretical analysis explores how PR-MPNNs enhance expressive power, and we identify precise conditions under which they outperform purely randomized approaches. Empirically, we demonstrate that our approach effectively mitigates issues like over-squashing and under-reaching. In addition, on established real-world datasets, our method exhibits competitive or superior predictive performance compared to traditional MPNN models and recent graph transformer architectures.

研究の動機と目的

• メッセージパッシングGNNの限界、特に固定されたグラフ構造に起因する過剰圧縮および表現力の制限を是正すること。
• 予測タスクに適合した、原理的で微分可能なグラフ構造の変更手法を構築すること。
• データ分布に適応せず、根拠のないランダムなエッジの追加を行うヒューリスティックな再接続手法の欠点を克服すること。
• 確率的再接続がモデルの表現力をどのように向上させるかを理論的に分析し、ランダム再接続を上回る条件を同定すること。
• 実世界のデータセットにおいて、PR-MPNNがMPNNおよびグラフトランスフォーマーと比較して優れたまたは競争力のある性能を示すことを実証的に検証すること。

提案手法

• PR-MPNNsは、上流モデルを用いて候補エッジの確率分布を学習し、正確にk本のエッジを持つkサブセット分布としてパrameter化する。
• この学習済み分布から複数のkエッジの隣接行列をサンプリングし、それらを集約（例：差分）することで最終的な微分可能なグラフ構造を形成する。
• バックプロパゲーションでは、kサブセット分布の正確で微分可能な周辺分布を用いることで、学習におけるバイアスと分散を低減する。
• 最近の正確で微分可能なkサブセットサンプリングの進展を活用し、エッジ選択メカニズムのエンドツーエンド学習を可能にする。
• 最終的なグラフ構造を下流のMPNNの入力として用い、構造と表現学習の両方を共同最適化可能にする。
• スケーラビリティと微分可能性を設計に組み込み、確率的エッジ選択プロセスを効果的にバックプロパゲーション可能にする。

実験結果

リサーチクエスチョン

• RQ1 固定グラフまたはランダムに再接続されたグラフと比較して、確率的再接続がメッセージパッシングGNNの表現力をどのように向上させるか、どのような条件下でその効果が顕著になるか。
• RQ2 微分可能なkサブセットサンプリング機構が、エッジ選択ポリシーの効果的かつ低分散な学習をどのように可能にするか。
• RQ3 PR-MPNNは、グラフニューラルネットワークにおける過剰圧縮および情報伝達の不足をどの程度緩和できるか。
• RQ4 実世界のグラフベンチマークにおいて、PR-MPNNはグラフトランスフォーマーや最先端のMPNNと比較して、同等または優れた性能を達成できるか。
• RQ5 学習されたエッジ選択ポリシーは、データ分布と予測タスクにどのように適合するか。

主な発見

• LRGBベンチマークのPeptides-structデータセットにおいて、PR-MPNNはMAEスコア0.2477 ± 0.0005を達成し、DRew-GIN (0.2799 ± 0.0016) や他のベースラインを上回った。
• 異種性が強いおよびトランスダクティブなWebKBデータセットにおいて、PR-MPNNはCornell (0.659 ± 0.040)、Texas (0.827 ± 0.032)、Wisconsin (0.750 ± 0.015) で最高の精度を記録し、すべてのMPNNベースラインを上回り、グラフトランスフォーマーと競合した。
• 長距離エッジの追加を学習することで、情報伝達の改善が図られ、過剰圧縮および情報伝達の不足が顕著に軽減された。これは、グローバル構造を必要とするタスクで向上した性能によって裏付けられた。
• ギャップ性質予測タスクにおいて、ベースラインGNNモデルと比較して、平均絶対誤差(MAE)が1.98倍改善された（1.70 ± 0.01 vs. 3.37 ± 0.03）。
• ZPVE性質においては、ベースラインGNNモデルと比較して14.13倍の改善が得られ、情報圧縮の緩和が強く示された（4.73 ± 0.08 vs. 66.87 ± 1.45）。
• 理論的分析により、学習されたエッジの分布とタスク関連構造との整合性が、正確な条件下で標準MPNNよりも高い表現力をPR-MPNNsが達成できることを確認した。