[論文レビュー] A Generalization of ViT/MLP-Mixer to Graphs

Xiaoxin He, Bryan Hooi+4 more|arXiv (Cornell University)|Dec 27, 2022

Advanced Graph Neural Networks被引用数 7

ひとこと要約

本稿では、グラフをパッチに分割し、トークンミキシングおよびチャネルミキシング機構を適用することで、ビジョントランスフォーマーやMLP-Mixerをグラフに一般化する新しいGNNアーキテクチャ「Graph ViT/MLP-Mixer」を提案する。ノード数とエッジ数に対して線形の計算量を達成し、過剰圧縮を軽減し、長距離依存性および表現力のベンチマークにおいてMP-GNNやグラフトランスフォーマーを上回る性能を発揮する。

ABSTRACT

Graph Neural Networks (GNNs) have shown great potential in the field of graph representation learning. Standard GNNs define a local message-passing mechanism which propagates information over the whole graph domain by stacking multiple layers. This paradigm suffers from two major limitations, over-squashing and poor long-range dependencies, that can be solved using global attention but significantly increases the computational cost to quadratic complexity. In this work, we propose an alternative approach to overcome these structural limitations by leveraging the ViT/MLP-Mixer architectures introduced in computer vision. We introduce a new class of GNNs, called Graph ViT/MLP-Mixer, that holds three key properties. First, they capture long-range dependency and mitigate the issue of over-squashing as demonstrated on Long Range Graph Benchmark and TreeNeighbourMatch datasets. Second, they offer better speed and memory efficiency with a complexity linear to the number of nodes and edges, surpassing the related Graph Transformer and expressive GNN models. Third, they show high expressivity in terms of graph isomorphism as they can distinguish at least 3-WL non-isomorphic graphs. We test our architecture on 4 simulated datasets and 7 real-world benchmarks, and show highly competitive results on all of them. The source code is available for reproducibility at: \url{https://github.com/XiaoxinHe/Graph-ViT-MLPMixer}.

研究の動機と目的

• メッセージパッシングGNNの限界、特に過剰圧縮と長距離依存性の低さを解消すること。
• グラフトランスフォーマーの二次的計算量を克服しつつ、表現力を維持または向上させること。
• 積み重ねられたメッセージパッシング層に依存せずに、長距離依存性を効率的に捉えるGNNアーキテクチャを開発すること。
• 1-WLを超えるグラフ同型性の表現力を高め、より複雑な非同型グラフを区別できるようにすること。
• 多様なグラフベンチマーク、特に大規模かつスパースなグラフにおいて、線形時間およびメモリ計算量で高い性能を達成すること。

提案手法

• グラフをMETISグラフ分割法を用いてP個のパッチに分割し、接続性を保つために1ホップの重複を設ける。
• 各パッチに対してベースとなるGNN（例：GCN、GIN）を適用し、パッチ埋め込みを計算する。これにより、標準のGNNに定数オーダーのオーバーヘッドが追加される。
• パッチ間の長距離依存性をモデル化するため、トランスフォーマー風のアテンションまたはフィードフォワード演算を実行するトークンミキサーレイヤーを適用する。
• 各パッチ内のノード表現を精緻化するため、チャネルミキサーレイヤーを適用し、特徴量レベルの相互作用を可能にする。
• 下流タスクのためのパッチ表現を統合するため、グローバル分類ヘッドを用いる。
• MLP-Mixer風のアーキテクチャを活用し、空間的（パッチ）およびチャネルワイドな特徴量学習を分離することで、効率的なグローバルモデリングを実現する。

実験結果

リサーチクエスチョン

• RQ1 ViT/MLP-Mixerにインspiredされたアーキテクチャを、グラフ構造データに効果的に一般化できるか？
• RQ2 提案されたGraph ViT/MLP-Mixerは、標準のMP-GNNやグラフトランスフォーマーに比べ、より優れた長距離依存性モデリングを達成できるか？
• RQ3 1-WLおよび2-WL GNNよりも高い表現力を維持しつつ、線形時間およびメモリ計算量を保ち続けられるか？
• RQ4 TreeNeighbourMatch や Long Range Graph Benchmark のような、長距離依存性および構造的複雑性を有するベンチマークで、モデルの性能はいかがなものか？
• RQ5 モデルの性能は、パッチ数やグラフ分割戦略にどの程度依存するか？

主な発見

• Peptides-funcで平均正解率に平均0.056の向上を達成し、Peptides-structではMAEが0.028低下し、優れた長距離相互作用モデリングを示した。
• TreeNeighbourMatchデータセットでは、深さr=7まで一般化が可能である一方、標準のMP-GNNはr=4で失敗しており、過剰圧縮の効果的な緩和が示された。
• SUNと同等の性能（Peptides-funcで0.6730、Peptides-structで0.2498）を達成したが、メモリ使用量は44倍少なく、学習速度は19倍速く、効率性が確認された。
• Graph MLP-Mixerは高い表現力を示し、少なくとも3-WL非同型グラフを区別できる能力を有しており、標準のMP-GNNを上回っている。
• モデルは時間およびメモリ計算量の両方で線形計算量O(N + E)を維持しており、スケーラビリティにおいてグラフトランスフォーマーやSUNのような表現力の高いGNNを大きく上回っている。
• 4つのシミュレート済みおよび7つの実世界ベンチマークにおける実験結果から、すべてのタスクで非常に競争力のある性能を示し、アーキテクチャの頑健性と一般化能力を裏付けた。