Graph-Mamba: Towards Long-Range Graph Sequence Modeling with Selective State Spaces

要約

アテンション メカニズムは、グラフ トランスフォーマーのノード間の長距離依存関係を把握するために広く使用されています。
二次計算コストがボトルネックとなり、アテンション メカニズムは大きなグラフでは拡張できません。
最近の計算効率の向上は主に、ランダムまたはヒューリスティックベースのグラフサブサンプリングによる注意の分散化によって達成されていますが、これではデータに依存したコンテキスト推論では不十分です。
Mamba などの状態空間モデル (SSM) は、シーケンシャル データの長距離依存関係をモデル化する際の有効性と効率性で注目を集めています。
ただし、SSM を非順次グラフ データに適応させるには、顕著な課題が生じます。
この研究では、Mamba ブロックを入力依存のノード選択メカニズムと統合することで、グラフ ネットワークにおける長距離コンテキスト モデリングを強化する最初の試みである Graph-Mamba を紹介します。
具体的には、グラフ中心のノードの優先順位付けと置換戦略を策定して、コンテキストを認識した推論を強化し、予測パフォーマンスの大幅な向上につながります。
10 のベンチマーク データセットでの広範な実験により、Graph-Mamba は、FLOP と GPU メモリ消費量の両方において数分の 1 の計算コストで、長距離グラフ予測タスクにおいて最先端の手法よりも優れたパフォーマンスを発揮することが実証されました。
コードとモデルは https://github.com/bowang-lab/Graph-Mamba で公開されています。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google