[論文レビュー] DOTIN: Dropping Task-Irrelevant Nodes for GNNs

Shaofeng Zhang, Feng Zhu+3 more|arXiv (Cornell University)|Apr 28, 2022

Advanced Graph Neural Networks被引用数 4

ひとこと要約

DOTIN は、タスク固有のグローバル表現を捉えるために K 個の学習可能な仮想ノードを活用することで、GNN におけるタスクに無関係なノードを最大 90% まで適応的に削除する、革新的なタスクに配慮したグラフプルーニング手法を提案する。トランスフォーマーに基づくアテンションメカニズムを用いてノードの重要度をランク付けすることで、D&D においてはパフォーマンスに損なわれることなく、約 50% の高速化と約 60% のメモリ使用量削減を達成し、単一およびマルチタスク設定の両方でプーリングベースのベースラインを上回る性能を発揮する。

ABSTRACT

Scalability is an important consideration for deep graph neural networks. Inspired by the conventional pooling layers in CNNs, many recent graph learning approaches have introduced the pooling strategy to reduce the size of graphs for learning, such that the scalability and efficiency can be improved. However, these pooling-based methods are mainly tailored to a single graph-level task and pay more attention to local information, limiting their performance in multi-task settings which often require task-specific global information. In this paper, departure from these pooling-based efforts, we design a new approach called DOTIN (\underline{D}r\underline{o}pping \underline{T}ask-\underline{I}rrelevant \underline{N}odes) to reduce the size of graphs. Specifically, by introducing $K$ learnable virtual nodes to represent the graph embeddings targeted to $K$ different graph-level tasks, respectively, up to 90\% raw nodes with low attentiveness with an attention model -- a transformer in this paper, can be adaptively dropped without notable performance decreasing. Achieving almost the same accuracy, our method speeds up GAT by about 50\% on graph-level tasks including graph classification and graph edit distance (GED) with about 60\% less memory, on D\&D dataset. Code will be made publicly available in https://github.com/Sherrylone/DOTIN.

研究の動機と目的

• GNN のスケーラビリティ制限を克服し、グラフレベルのタスクにおける計算コストとメモリコストを低減すること。
• 局所的情報を優先し、タスクに依存しない既存のプーリング手法の欠点を是正すること。
• 仮想ノードを介してグローバル表現を学習することで、タスク固有の適応的ノードプルーニングを可能にすること。
• パフォーマンスを損なわず、マルチタスク GNN 環境における効率性とスケーラビリティを向上させること。
• GNN に依存しない、パラメータ効率の良い階層的グラフサイズ縮小手法を提供すること。

提案手法

• K 個の学習可能な仮想ノードを導入し、K 個の異なるグラフレベルタスクごとにタスク固有のグローバル埋め込みを表す。
• すべてのタスクにわたるノードの注意スコアを計算するために、トランスフォーマーに基づくアテンションメカニズムを用いる。
• タスク間で平均的な注意スコアが低いノード（＝タスクに無関係なノード）を特定し、削除する。
• 各 GNN 層の後に階層的にノード削除を適用することで、追加の GNN 层やクラスタリング層を必要とせず、グラフサイズを縮小する。
• 特にマルチタスク環境において、タスク固有の注意スコアが高くなるノードを保持することで、モデルのパフォーマンスを維持する。
• 仮想ノードの注意スコアをソートすることで O(N log N) の計算量を達成し、高価なクラスタリングや追加の GNN 層を回避する。

実験結果

リサーチクエスチョン

• RQ1 グラフレベルのタスクにおいて、GNN のパフォーマンスが低下することなく、タスクに無関係なノードを特定・削除することは可能か？
• RQ2 分類やグラフ編集距離（GED）のような異なるグラフレベルタスクにおいて、タスク固有のノード重要度はどのように変化するか？
• RQ3 GNN に依存しない手法が、プーリングベースのアプローチよりも優れた効率性とスケーラビリティを達成できるか？
• RQ4 マルチタスク注意に基づく適応的ノード削除は、パフォーマンスの向上とメモリ使用量の削減をもたらすか？
• RQ5 本手法は、未観測のタスクやデータセットに対しても一般化可能であり、効率性と正確性を維持できるか？

主な発見

• D&D データセットでは、ベースラインと比較して顕著なパフォーマンス低下が見られず、ノード数を最大 90% 削減した。
• D&D データセットにおいて、GAT バックボーンを用いた場合、約 50% の高速化と約 60% のメモリ使用量削減を達成した。
• 隠れ層次元が 512 および 4096 の場合、90% のノード削除率を適用した DOTIN は、それぞれベースラインと比較して 42% および 57% のメモリ使用量に抑えた。
• 可視化結果から、ノードの重要度はタスクによって顕著に異なることが示された — 例えば、グラフ分類において重要なノードは GED ではしばしば無関係であり、逆も同様である。
• gPool などのプーリングベースの手法とは異なり、タスクに依存しない設計のため、タスクに必要なノードを誤って削除してしまう可能性があるが、DOTIN はその問題を回避している。
• 複数のデータセット（D&D, PROTEINS, NCI1, NCI109, FRANKENSEIN）にわたり、さまざまなバックボーンやノード削除率においても、高い正確性を維持し、一般化性能に優れている。