[論文レビュー] Structure-Aware Transformer for Graph Representation Learning

Dexiong Chen, Leslie O’Bray+1 more|arXiv (Cornell University)|Feb 7, 2022

Advanced Graph Neural Networks被引用数 46

ひとこと要約

Structure-Aware Transformer (SAT) は graph transformers を構造認識型自己注意機構で強化し、局所サブグラフ情報を取り込むことで、GNNベースのサブグラフ抽出と Transformer attention を組み合わせることにより、いくつかのグラフベンチマークで最先端の結果を達成します。

ABSTRACT

The Transformer architecture has gained growing attention in graph representation learning recently, as it naturally overcomes several limitations of graph neural networks (GNNs) by avoiding their strict structural inductive biases and instead only encoding the graph structure via positional encoding. Here, we show that the node representations generated by the Transformer with positional encoding do not necessarily capture structural similarity between them. To address this issue, we propose the Structure-Aware Transformer, a class of simple and flexible graph Transformers built upon a new self-attention mechanism. This new self-attention incorporates structural information into the original self-attention by extracting a subgraph representation rooted at each node before computing the attention. We propose several methods for automatically generating the subgraph representation and show theoretically that the resulting representations are at least as expressive as the subgraph representations. Empirically, our method achieves state-of-the-art performance on five graph prediction benchmarks. Our structure-aware framework can leverage any existing GNN to extract the subgraph representation, and we show that it systematically improves performance relative to the base GNN model, successfully combining the advantages of GNNs and Transformers. Our code is available at https://github.com/BorgwardtLab/SAT.

研究の動機と目的

• グラフ表現学習における標準的な GNN の限界（表現力、過度平滑化、過度圧縮）を解決する。
• 属性ベースの注意だけに依存するのを超え、明示的な構造情報を Transformer の注意機構に埋め込む。
• 既存の GNN を再利用またはプラグインしてサブグラフ表現を抽出し、全体的な性能を向上させる。
• 表現力に関する理論的保証を提供し、さまざまなグラフタスクで経験的な向上を示す。

提案手法

• Transformer の自己注意をカーネルスムーザーとして再定式化し、局所的なグラフ構造をサブグラフを介して考慮するように指数カーネルを拡張する。
• Introduce SA-attn, a structure-aware attention that uses a subgraph representation S_G(v) for each node and a kernel κ_graph comparing subgraphs: SA-attn(v) = sum_u κ_graph(S_G(v), S_G(u)) / sum_w κ_graph(S_G(v), S_G(w)) f(x_u).
• 構造抽出器 φ(u,G) を定義してサブグラフ表現を生成し、k-subtree GNN 抽出器および k-subgraph GNN 抽出器を含め、元のノード特徴と連結することもできる。
• 構造抽出器を任意の微分可能なモデル（GNN、グラフカーネル）として許可し、適切な抽出器を介してエッジ属性をサポートする。
• スキップ接続、FFN、レイヤ norm、そして高度に接続されたノードを緩和するための次数ベースのスキップ因子を備えた Transformer ブロックに構造認識型注意を組み込む。
• SAT を絶対エンコーディング（e.g., RWPE）と組み合わせて補完的な情報を取り入れ、さらなる性能向上を図る。

実験結果

リサーチクエスチョン

• RQ1 Explicit structure-aware self-attention は絶対的位置エンコーディングが提供する以上のノード間の構造的類似性を捉えられるか。
• RQ2 異なるサブグラフ抽出器（k-subtree vs k-subgraph、基盤 GNN の選択）の予測性能への影響はどうなるか。
• RQ3 SAT は構造抽出器が用いるサブグラフ表現に対して表現力の理論的保証を提供するか。
• RQ4 SAT はグラフおよびノード予測ベンチマークで最先端の GNN とグラフトランスフォーマーと比較してどうか。
• RQ5 SAT は構造認識型注意を提供することで、どんな基盤となる GNN でも実用的に拡張できるか。

主な発見

• SAT は five graph prediction benchmarks で最先端の性能を達成し、GNN とグラフ Transformers の双方を上回った。
• 標準の自己注意を構造認識型 SA-attn に置換または補足することは、表現力が基になるサブグラフ表現と少なくとも同等になる表現を生み出す。
• Both k-subtree SAT and k-subgraph SAT の両方が複数のデータセットで基盤となる GNN を一貫して改善し、k-subgraph の方が高い表現力を示すことが多い。
• Incorporating structural information via SAT provides notable gains over vanilla Transformer with RWPE absolute encoding alone.
• A theoretical result shows SA-attn preserves expressiveness relative to the subgraph extractor, and a Lipschitz-based bound relates node representation similarity to subgraph and feature similarity (Theorem 1).
• Experimental results on OGB datasets (CODE2, PPA) illustrate strong performance gains, with SAT variants outperforming several baselines.