Graph Mamba: Towards Learning on Graphs with State Space Models

要約

グラフ ニューラル ネットワーク (GNN) は、グラフ表現の学習において有望な可能性を示しています。
大多数の GNN はローカルのメッセージ受け渡しメカニズムを定義し、複数のレイヤーを積み重ねることによってグラフ上に情報を伝播します。
ただし、これらの方法には 2 つの大きな制限があることが知られています。それは、過剰な潰しと、長距離の依存関係の捕捉が不十分であるということです。
最近、メッセージ パッシング ニューラル ネットワーク (MPNN) の強力な代替手段として、グラフ トランスフォーマー (GT) が登場しました。
ただし、GT には 2 次の計算コストがあり、グラフ構造に対する誘導バイアスがなく、複雑な位置/構造エンコーディング (SE/PE) に依存します。
この論文では、実際にはトランスフォーマー、複雑なメッセージパッシング、および SE/PE が優れたパフォーマンスを実現するのに十分である一方で、どちらも必要ではないことを示します。
Mamba などの状態空間モデル (SSM) の最近の成功を動機として、選択的 SSM に基づく新しいクラスの GNN の一般的なフレームワークである Graph Mamba Networks (GMN) を紹介します。
SSM をグラフ構造データに適応させる際の新しい課題を議論して分類し、GMN を設計するための 4 つの必須ステップと 1 つのオプションのステップを示します。ここでは、(1) 近隣トークン化、(2) トークン順序付け、(3) 双方向選択のアーキテクチャーを選択します。
SSM エンコーダ、(4) ローカル エンコーディング、および必須の (5) PE および SE。
さらに、GMN の力の理論的根拠を提供します。
実験により、GMN は計算コストがはるかに低いにもかかわらず、長距離、小規模、大規模、および異好性のベンチマーク データセットにおいて優れたパフォーマンスを達成することが実証されています。

要約(オリジナル)

Graph Neural Networks (GNNs) have shown promising potential in graph representation learning. The majority of GNNs define a local message-passing mechanism, propagating information over the graph by stacking multiple layers. These methods, however, are known to suffer from two major limitations: over-squashing and poor capturing of long-range dependencies. Recently, Graph Transformers (GTs) emerged as a powerful alternative to Message-Passing Neural Networks (MPNNs). GTs, however, have quadratic computational cost, lack inductive biases on graph structures, and rely on complex Positional/Structural Encodings (SE/PE). In this paper, we show that while Transformers, complex message-passing, and SE/PE are sufficient for good performance in practice, neither is necessary. Motivated by the recent success of State Space Models (SSMs), such as Mamba, we present Graph Mamba Networks (GMNs), a general framework for a new class of GNNs based on selective SSMs. We discuss and categorize the new challenges when adapting SSMs to graph-structured data, and present four required and one optional steps to design GMNs, where we choose (1) Neighborhood Tokenization, (2) Token Ordering, (3) Architecture of Bidirectional Selective SSM Encoder, (4) Local Encoding, and dispensable (5) PE and SE. We further provide theoretical justification for the power of GMNs. Experiments demonstrate that despite much less computational cost, GMNs attain an outstanding performance in long-range, small-scale, large-scale, and heterophilic benchmark datasets.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google