[論文レビュー] Graph-Mamba: Towards Long-Range Graph Sequence Modeling with Selective State Spaces

Xiaoming Wang, Oleksii Tsepa+2 more|arXiv (Cornell University)|Feb 1, 2024

Graph Theory and Algorithms被引用数 20

ひとこと要約

Graph-Mamba は GraphGPS の注意機構を置換する選択的状態空間モデル（SSM）ベースのグラフブロックを導入し、データ依存の長距離文脈モデリングを線形時間計算量で実現し、メモリを削減します。密な注意機構および疎な注意機構と比べて、長距離グラフベンチマークで優れた性能を示します。

ABSTRACT

Attention mechanisms have been widely used to capture long-range dependencies among nodes in Graph Transformers. Bottlenecked by the quadratic computational cost, attention mechanisms fail to scale in large graphs. Recent improvements in computational efficiency are mainly achieved by attention sparsification with random or heuristic-based graph subsampling, which falls short in data-dependent context reasoning. State space models (SSMs), such as Mamba, have gained prominence for their effectiveness and efficiency in modeling long-range dependencies in sequential data. However, adapting SSMs to non-sequential graph data presents a notable challenge. In this work, we introduce Graph-Mamba, the first attempt to enhance long-range context modeling in graph networks by integrating a Mamba block with the input-dependent node selection mechanism. Specifically, we formulate graph-centric node prioritization and permutation strategies to enhance context-aware reasoning, leading to a substantial improvement in predictive performance. Extensive experiments on ten benchmark datasets demonstrate that Graph-Mamba outperforms state-of-the-art methods in long-range graph prediction tasks, with a fraction of the computational cost in both FLOPs and GPU memory consumption. The code and models are publicly available at https://github.com/bowang-lab/Graph-Mamba.

研究の動機と目的

• 大規模グラフにおける長距離依存性のための Graph Transformers の計算ボトルネックを動機づけ、対処する。
• 入力依存ノード選択と文脈フィルタリングを行う、グラフ適応型選択的状態空間モデリング（GMB）モジュールを開発する。
• GMB を GraphGPS フレームワークに統合して Graph-Mamba を作成し、非連結グラフデータ上での効率的な長距離推論を可能にする。
• グラフ特有の系列モデリングを安定化・改善するためのノード優先度付けと順列ベースの訓練/推論レシピを提案する。

提案手法

• 選択的状態空間モデルである Mamba を、グラフ中心のノード優先度付けと順列戦略を導入してグラフへ適用する。
• グラフGPS の注意機構を、局所文脈を MPNN ベースで組み合わせ、グローバル文脈のためのノード中心 GMB を備えた Graph-Mamba ブロック（GMB）に置き換える。
• 入力依存ノード選択と系列認識可能なノード優先度付けという、2段階のスパース化を可能にするデータ依存選択メカニズムを実装する。
• 入力長に対して線形時間計算量を達成するために、GPU に適した状態空間更新の離散化/パラメータ化を用いる。
• ノード順序間での置換不変性と安定性を促進するため、順列ベースの訓練と推論レシピを導入する。
• Evaluates on ten Long Range Graph Benchmark データセットで評価し、密集型 Transformer および疎注意のベースラインと比較する。）

実験結果

リサーチクエスチョン

• RQ1 グラフ適応型選択的状態空間モデル（Graph-Mamba）は、長距離依存性タスクにおいて、計算コストを削減しつつ、密集/疎のグラフ注意機構と同等またはそれ以上の予測力を発揮できるか？
• RQ2 入力依存ノード優先度付けと順列ベースの訓練は、非連結グラフデータにおける長距離推論を改善するか？
• RQ3 最先端の注意機構と比較して、Graph-Mamba は多様なグラフ予測タスク（ノード、グラフ、エッジレベル）およびグラフサイズにおいてどのように性能を示すか？

主な発見

モデル	ペプチド-機能 AP↑	ペプチド-構造 MAE↓	PascalVOC-SP F1↑	COCO-SP F1↑	MALNET-TINY Accuracy↑
GCN	0.5930 ± 0.0023	0.3496 ± 0.0013	0.1268 ± 0.0060	0.0841 ± 0.0010	0.8100
GIN	0.5498 ± 0.0079	0.3547 ± 0.0045	0.1265 ± 0.0076	0.1339 ± 0.0044	0.8898 ± 0.0055
GatedGCN	0.5864 ± 0.0077	0.3420 ± 0.0013	0.2873 ± 0.0219	0.2641 ± 0.0045	0.9223 ± 0.0065
GPS+Transformer	0.6575 ± 0.0049	0.2510 ± 0.0015	0.3689 ± 0.0131	0.3774 ± 0.0150	OOM (bs=8)
GPS+Performer	0.6475 ± 0.0056	0.2558 ± 0.0012	0.3724 ± 0.0131	0.3761 ± 0.0101	0.9264 ± 0.0078
GPS+BigBird	0.5854 ± 0.0079	0.2842 ± 0.0130	0.2762 ± 0.0069	0.2622 ± 0.0008	0.9234 ± 0.0034
Exphormer	0.6258 ± 0.0092	0.2512 ± 0.0025	0.3446 ± 0.0064	0.3430 ± 0.0108	0.9422 ± 0.0024
Graph-Mamba	0.6739 ± 0.0087	0.2478 ± 0.0016	0.4191 ± 0.0126	0.3960 ± 0.0175	0.9340 ± 0.0027

• Graph-Mamba は、10データセットにわたる長距離グラフ予測タスクで最先端または競争力のある結果を達成し、特にノード数が150–1,400のグラフで顕著。
• 大規模グラフでは、他の疎注意法より最大で5%の改善を提供し、いくつかのケースで密な Transformer ベースラインを上回る。
• Graph-Mamba は線形時間計算量と GPU メモリ使用量の削減を示し、大規模グラフで最大74%のメモリ削減、MalNet-Tiny では Transformer 比で FLOPsを66%削減。
• 度に基づくノード優先度付けとノードレベルの順列を組み合わせることで顕著な性能向上を得られ、デフォルトの訓練/推論レシピとして推奨される。
• 実測 FLOPs/メモリのベンチマークは、Graph-Mamba がデータセット全体で密なおよび疎な注意ベースラインよりコストが低いことを示し、より大きなグラフ上での訓練を可能にする。）