[論文レビュー] Stochastic Aggregation in Graph Neural Networks

Yuanqing Wang, Theofanis Karaletsos|arXiv (Cornell University)|Feb 25, 2021

Advanced Graph Neural Networks参考文献 24被引用数 7

ひとこと要約

本稿では、グラフニューラルネットワーク（GNN）におけるメッセージパッシング中にエッジ重みに適応的ノイズを注入する統合的フレームワーク「Stochastic Aggregation（STAG）」を提案する。これにより表現力が向上し、過剰平滑化が軽減される。STAGは、学術論文および分子グラフベンチマークにおいて性能を向上させ、変分推論を用いたバージョンは最小限の計算コストで最先端の結果を達成する。

ABSTRACT

Graph neural networks (GNNs) manifest pathologies including over-smoothing and limited discriminating power as a result of suboptimally expressive aggregating mechanisms. We herein present a unifying framework for stochastic aggregation (STAG) in GNNs, where noise is (adaptively) injected into the aggregation process from the neighborhood to form node embeddings. We provide theoretical arguments that STAG models, with little overhead, remedy both of the aforementioned problems. In addition to fixed-noise models, we also propose probabilistic versions of STAG models and a variational inference framework to learn the noise posterior. We conduct illustrative experiments clearly targeting oversmoothing and multiset aggregation limitations. Furthermore, STAG enhances general performance of GNNs demonstrated by competitive performance in common citation and molecule graph benchmark datasets.

研究の動機と目的

• 決定的集約メカニズムに起因するGNNにおける過剰平滑化と表現力の制限を解消すること。
• ドロップアウト、ドロップエッジ、グラフドロップコネクトといった既存の正則化手法を、統一的な確率的フレームワークに統合すること。
• ノイズ分布パラメータをモデルパラメータと共に同時に学習する変分推論（VI）に基づくSTAGの変種を開発すること。
• 標準的なGNNベンチマーク（学術論文および分子グラフを含む）において、STAGの有効性を実証的に検証すること。
• エッジ重みにおける確率性が、顕著な計算コスト増加なしに一般化性能とモデルの頑健性を向上させることを示すこと。

提案手法

• STAGは、各メッセージパッシングステップにおいてノイズ分布からの確率的エッジ重みサンプリングに、決定的集約を置き換える。
• ノイズは、連続分布（例：正規分布N(1,1)）または離散分布（例：ベルヌーイ分布）を用いてエッジ重みを摂動させることで注入され、近隣メッセージの再重み付けが可能になる。
• 学習では、モンテカルロサンプリングによる勾配推定が行われ、バイアスなしにエンドツーエンドのバックプロパゲーションが可能になる。
• 変分推論（VI）フレームワークを導入し、ノイズ分布のパラメータを変分パラメータとして学習する。これは、ノード特徴量およびグラフ構造に条件付けられる可能性がある。
• 推論時の予測後確率分布は、学習済みエッジ重み分布の周辺化によって得られる。
• PyTorchおよびDGLにおいて効率的に実装されており、メッセージパッシング操作への変更は1行で完了する。

実験結果

リサーチクエスチョン

• RQ1 GNNにおけるエッジ重みの確率的摂動は、決定的集約と比較してモデルの表現力向上と過剰平滑化の軽減に寄与するか？
• RQ2 ドロップアウト、ドロップエッジ、グラフドロップコネクトといった既存の正則化手法と比較して、STAGの性能と一般化性能はどのように異なるか？
• RQ3 ノイズパラメータの学習に変分推論フレームワークを用いることで、多様なグラフ構造タスクにおけるGNN性能向上が達成可能か？
• RQ4 特に正規化が適用されていない場合に、連続的ノイズ分布（例：正規分布）が離散的ノイズ分布（例：ベルヌーイ分布）を上回る性能を示すか？
• RQ5 STAGは、学術ネットワークや分子グラフを含むさまざまなグラフタイプに一般化可能で、一貫した性能向上を示すか？

主な発見

• 連続的ノイズ分布（例：正規分布N(1,1)）を用いたSTAGは、CoraおよびCiteseerにおいて決定的ベースラインおよび離散的ノイズバージョンを上回り、複数回の実験で一貫した向上を示す。
• 変分推論を用いたSTAG_VIは、学術および分子グラフベンチマークの両方で非適応的STAGおよびグラフドロップコネクトを常に上回る性能を達成する。
• エッジおよび特徴量に依存するノイズパラメータを学習する最も表現力のあるSTAG_VIバージョンは、ESOLおよびFreeSolvデータセットで最先端の性能を達成する。
• V100 GPU上では、STAGによる1反復の順伝播処理で、計算時間は5.9〜9.3 msの増加にとどまり、その軽量性が裏付けられる。
• 正規化操作は連続的ノイズ分布を用いる場合に性能を低下させることが示され、STAGのノイズ注入が、このような補正なしに本質的に安定であることが示唆される。
• 本フレームワークは、社会的ネットワークおよび分子グラフデータセットの両方で、一般化性能が高く、一貫した向上を示す。