[論文レビュー] MAG-GNN: Reinforcement Learning Boosted Graph Neural Network

Lecheng Kong, Jiarui Feng+4 more|arXiv (Cornell University)|Oct 29, 2023

Advanced Graph Neural Networks被引用数 5

ひとこと要約

MAG-GNN は、グラフニューラルネットワーク（GNN）のための最小で表現力のある部分グラフのサブセットを特定する強化学習（RL）ベースの手法を提案する。これは、すべての部分グラフを列挙する計算コストの高いプロセスに代わるものである。反復的に識別力の高い部分グラフを選択する RL エージェントを用いることで、MAG-GNN は複数のグラフベンチマークで最先端の性能を達成するとともに、部分グラフ GNN よりも顕著に推論時間を短縮した。

ABSTRACT

While Graph Neural Networks (GNNs) recently became powerful tools in graph learning tasks, considerable efforts have been spent on improving GNNs' structural encoding ability. A particular line of work proposed subgraph GNNs that use subgraph information to improve GNNs' expressivity and achieved great success. However, such effectivity sacrifices the efficiency of GNNs by enumerating all possible subgraphs. In this paper, we analyze the necessity of complete subgraph enumeration and show that a model can achieve a comparable level of expressivity by considering a small subset of the subgraphs. We then formulate the identification of the optimal subset as a combinatorial optimization problem and propose Magnetic Graph Neural Network (MAG-GNN), a reinforcement learning (RL) boosted GNN, to solve the problem. Starting with a candidate subgraph set, MAG-GNN employs an RL agent to iteratively update the subgraphs to locate the most expressive set for prediction. This reduces the exponential complexity of subgraph enumeration to the constant complexity of a subgraph search algorithm while keeping good expressivity. We conduct extensive experiments on many datasets, showing that MAG-GNN achieves competitive performance to state-of-the-art methods and even outperforms many subgraph GNNs. We also demonstrate that MAG-GNN effectively reduces the running time of subgraph GNNs.

研究の動機と目的

• 部分グラフ GNN が表現力を向上させるためにすべての可能な部分グラフを列挙することに伴う高い計算コストに対処すること。
• 全部分グラフ列挙と同等の表現力を達成できる、小さな最適な部分グラフサブセットが存在するかどうかを調査すること。
• 全列挙を伴わずに高い GNN の表現力を維持できる、効率的でスケーラブルな手法を開発すること。
• 部分グラフ GNN の推論時間を短縮しつつ、予測性能を維持または向上させること。

提案手法

• MAG-GNN は、入力グラフからランダムにサンプリングされたルート付き部分グラフの候補集合を出発点とする。
• RL エージェントは、モデル性能に基づく報酬信号を最大化するように、反復的に候補集合内の部分グラフを置き換えるように訓練される。
• エージェントは、ターゲット部分グラフを新しい候補部分グラフに置き換える際の期待報酬を推定するために Q テーブルを用いる。
• 部分グラフの置き換えプロセスは、各ステップで期待報酬が最大となる更新を選び出す方策に従う。
• 最終的に選ばれた部分グラフセットは、下流の GNN に入力として与えられ、予測タスクに使用される。
• この手法により、指数的になる部分グラフ列挙の計算量が定数ステップの RL 探索に低減され、効率的で表現力のある GNN の実現が可能になった。

実験結果

リサーチクエスチョン

• RQ1 小さな最適選択部分グラフサブセットは、GNN における全部分グラフ列挙と同等の表現力を達成できるか？
• RQ2 強化学習は、GNN の性能向上に寄与する最も識別力の高い部分グラフを効果的に特定できるか？
• RQ3 部分グラフ列挙を定数ステップの RL プロセスに短縮することで、GNN の性能は維持または向上するか？ また、実行時間は短縮されるか？
• RQ4 MAG-GNN は、最先端の部分グラフ GNN と比較して、精度と効率の両面で優れているか？

主な発見

• MAG-GNN は最先端の手法と同等の性能を達成し、ZINC や OGBG-MOLHIV を含む複数のデータセットで、いくつかの部分グラフ GNN よりも優れた性能を示した。
• ZINC データセットでは、MAG-GNN はテスト損失 0.106 ± 0.014 を達成し、GIN（0.163 ± 0.004）を顕著に上回り、CIN（0.079 ± 0.006）のようなより複雑なモデルに近い性能に到達した。
• 事前学習（MAG-GNN-PRE）を適用すると、さらに性能が向上し、ZINC ではテスト損失 0.096 ± 0.009、OGBG-MOLHIV では 78.30 ± 1.08 を達成した。
• MAG-GNN は推論時間を顕著に短縮しており、I²GNN や NGNN よりも高速であり、類似した効率性を維持しながらも、精度で上回った。
• モデルはさまざまなベース GNN に適応可能であり、アーキテクチャ統合によるさらなる性能向上の可能性を示唆している。
• アブレーションスタディにより、MAG-GNN の RL ベースの部分グラフ選択が、固定構造の部分グラフ GNN で見られる過学習を避けるとともに、表現力と効率性のバランスを効果的に実現していることが確認された。