Graph Mamba: Towards Learning on Graphs with State Space Models

Graph Mamba: Towards Learning on Graphs with State Space Models

19 Feb 2024 | Ali Behrouz * 1 Farnoosh Hashemi * 1
Graph Neural Networks (GNNs) have shown promise in graph representation learning, but they suffer from limitations such as over-squashing and poor capture of long-range dependencies. Graph Transformers (GTs) offer a powerful alternative, but they have quadratic computational costs and rely on complex positional/structural encodings (PE/SE). This paper introduces Graph Mamba Networks (GMNs), a new framework for GNNs based on selective State Space Models (SSMs). GMNs address these limitations by using a more efficient tokenization process, a bidirectional SSM encoder, and optional PE/SE. The paper discusses the challenges of adapting SSMs to graph-structured data and presents a five-step recipe for designing GMNs. Theoretical justification and experimental results demonstrate that GMNs achieve outstanding performance on various benchmark datasets, including long-range, small-scale, large-scale, and heterophilic graphs, with reduced computational costs compared to GTs. The code for GMNs is available online.Graph Neural Networks (GNNs) have shown promise in graph representation learning, but they suffer from limitations such as over-squashing and poor capture of long-range dependencies. Graph Transformers (GTs) offer a powerful alternative, but they have quadratic computational costs and rely on complex positional/structural encodings (PE/SE). This paper introduces Graph Mamba Networks (GMNs), a new framework for GNNs based on selective State Space Models (SSMs). GMNs address these limitations by using a more efficient tokenization process, a bidirectional SSM encoder, and optional PE/SE. The paper discusses the challenges of adapting SSMs to graph-structured data and presents a five-step recipe for designing GMNs. Theoretical justification and experimental results demonstrate that GMNs achieve outstanding performance on various benchmark datasets, including long-range, small-scale, large-scale, and heterophilic graphs, with reduced computational costs compared to GTs. The code for GMNs is available online.
Reach us at info@study.space