[論文レビュー] NodeFormer: A Scalable Graph Structure Learning Transformer for Node Classification

Qitian Wu, Wentao Zhao+3 more|arXiv (Cornell University)|Jun 14, 2023

Advanced Graph Neural Networks被引用数 75

ひとこと要約

NodeFormer は、潜在グラフ構造を学習するためのカーネル化された Gumbel-Softmax を備えた全対のメッセージパッシングを行う Transformer を導入し、層ごとに線形時間計算量を実現し、ノード分類のために数百万ノード規模のグラフへスケーリングする。

ABSTRACT

Graph neural networks have been extensively studied for learning with inter-connected data. Despite this, recent evidence has revealed GNNs' deficiencies related to over-squashing, heterophily, handling long-range dependencies, edge incompleteness and particularly, the absence of graphs altogether. While a plausible solution is to learn new adaptive topology for message passing, issues concerning quadratic complexity hinder simultaneous guarantees for scalability and precision in large networks. In this paper, we introduce a novel all-pair message passing scheme for efficiently propagating node signals between arbitrary nodes, as an important building block for a pioneering Transformer-style network for node classification on large graphs, dubbed as extsc{NodeFormer}. Specifically, the efficient computation is enabled by a kernerlized Gumbel-Softmax operator that reduces the algorithmic complexity to linearity w.r.t. node numbers for learning latent graph structures from large, potentially fully-connected graphs in a differentiable manner. We also provide accompanying theory as justification for our design. Extensive experiments demonstrate the promising efficacy of the method in various tasks including node classification on graphs (with up to 2M nodes) and graph-enhanced applications (e.g., image classification) where input graphs are missing.

研究の動機と目的

• 固定された入力トポロジーを超えて、ノード分類のための柔軟な潜在グラフ構造の学習を動機付ける。
• すべてのノードペア間で潜在グラフを推定する、スケーラブルで微分可能な機構を開発する。
• 入力グラフを持たない大規模グラフに対して、Transformer風アーキテクチャをエンドツーエンドで訓練できるようにする。
• 利用可能な場合は、関係的バイアスとエッジ正則化を通じて入力トポロジーを取り入れる。
• 最大200万ノードまでのグラフに対するスケーラビリティと、ヘトロフィリックおよびグラフなし設定での有効性を示す。

提案手法

• 層ごとに学習される潜在グラフ上で全対のメッセージパッシングを行う NodeFormer を提案する。
• ランダム特徴写像によって全対アテンションを線形時間で近似するカーネル化された Gumbel-Softmax 演算子を用いる。
• ランダム特徴で近似されたカーネル関数を用いてアテンションを再表現し、N^2 の計算を回避する。
• 温度付きの Gumbel-Softmax を用いた微分可能なサンプリングスキームを適用して、疎で学習可能な潜在グラフを得る。
• 利用可能なときに学習されたエッジを観測エッジと一致させるエッジレベルの正則化損失を導入する。
• エッジが存在する場合、入力グラフからの関係的バイアスをアテンション重みに加法项として組み込む。

Figure 1: Illustration for the data flow of NodeFormer which takes node embedding matrix $\mathbf{X}$ and (optional) graph adjacency matrix $\mathbf{A}$ as input. There are three components in NodeFormer . The first one is the all-pair message passing (MP) module (colored red) which adopts our propo

実験結果

リサーチクエスチョン

• RQ1 入力トポロジーに依存せず、大規模グラフ上でメッセージパッシングのための効果的な潜在グラフを学習できる Transformer ライクなノード分類器はあるだろうか？
• RQ2 離散的な潜在構造に対して、微分可能な最適化を用いたスケーラブルな全対のメッセージパッシングをどのように実現できるか？
• RQ3 同質性・異質性のグラフにおける潜在構造学習がノード分類性能に与える影響はどのようなものか？
• RQ4 カーネル特徴写像の次元、温度、近似精度のトレードオフは、カーネル化された Gumbel-Softmax においてどのようになるか？
• RQ5 大規模データセット（ノード数百万レベル）に対して、NodeFormer は既存の構造学習法やスケーラブルな GNN 手法とどう比較されるか？

主な発見

モデル	ROC-AUC (%)	訓練メモリ
MLP	72.04 ± 0.48	2.0 GB
GCN	72.51 ± 0.35	2.5 GB
SGC	70.31 ± 0.23	1.2 GB
GraphSAINT-GCN	73.51 ± 1.31	2.3 GB
GraphSAINT-GAT	74.63 ± 1.24	5.2 GB
NodeFormer	77.45 ± 1.15	3.2 GB
NodeFormer -dt	75.50 ± 0.64	3.1 GB
NodeFormer -tp	76.18 ± 0.09	3.2 GB

• NodeFormer は、ヘトロフィリックなグラフを含む標準的なノード分類ベンチマークで競争力のある、または優れた精度/ROC-AUC を達成する。
• 大規模データセット（OGB-Proteins および Amazon2M）では、NodeFormer はベースラインより高い指標を達成しつつ、実現可能なメモリ使用量を維持する。
• 本手法は最大200万ノードのグラフへ線形の層あたり計算量でスケールし、速度とメモリの点で従来の構造学習法を上回る。
• ランダム特徴を用いたカーネル化された Gumbel-Softmax は、N に依存しない誤差を持つ適切な近似を提供し、スケーラブルな潜在構造学習を可能にする。
• 関係的バイアスとエッジ正則化は、利用可能な場合に既存の入力トポロジーを活用するのに役立ち、訓練の安定性と性能を向上させる。
• NodeFormer は入力グラフが欠如している場合にも、グラフ強化タスク（画像・テキスト分類）への適用性を示す。

Figure 2: Experimental results for node classification in transductive setting on four common datasets. The missing results on Deezer is caused by out-of-memory (OOM).