PaperlayerView exact revision on arXiv

Source-linked AI summary

Graph Mamba: Towards Learning on Graphs with State Space Models

Ali Behrouz, Farnoosh Hashemi

arXiv:2402.08678v2cs.LG

TL;DR

Graph learning needs to capture long-range dependencies without the over-squashing of local message passing or the quadratic cost of global attention. The paper proposes Graph Mamba Networks, a selective-SSM framework for graphs, and reports strong performance across diverse benchmarks with lower computational cost.

  • Problem

    Graph Transformers address long-range interactions but have O(n^2) computational cost, while adapting sequence-oriented Mamba directly to graphs is challenging.

  • Method

    GMNs combine neighborhood tokenization, token ordering, local encoding, and bidirectional selective SSM encoding, with PE/SE optional.

  • Results

    GMNs attain outstanding performance on long-range, small-scale, large-scale, and heterophilic benchmarks while consuming less GPU memory.

  • Takeaways & Limitations

    GMNs show that Transformers, complex message passing, and PE/SE are not necessary for strong graph-learning performance in the reported experiments.

  • Takeaways & Limitations

    The concurrent GPS-based Mamba variant remains bounded by its MPNN’s expressive power and cannot distinguish graphs sharing the same degree sequence in the stated worst case.

Abstract

from arXiv · show

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.

1. Introduction

Graph learning methods trade off local message passing, global attention, and sequence modeling, each with limitations in long-range modeling, scalability, or graph structure. The paper introduces Graph Mamba Networks as a selective-SSM framework with a structured design recipe and strong benchmark performance.

  • Motivation: MPNNs aggregate neighborhood information locally, while Graph Transformers model global interactions but incur O(n^2) time and memory costs.MPNNs are also associated with over-squashing and limited long-range dependency capture.
  • Motivation: Directly applying Mamba to graphs is challenging because its recurrent sequence processing does not naturally match complex, non-causal graph structure.Naive replacements of Transformers with Mamba in existing graph frameworks reportedly give suboptimal effectiveness and time efficiency.
  • Graph Mamba Networks: GMNs provide four required steps—tokenization, token ordering, local encoding, and bidirectional selective SSM encoding—with PE/SE as an optional step.The framework is presented as a general class of graph-learning models based on state space models.
  • Graph Mamba Networks: The proposed tokenization bridges node- and subgraph-level representations using one parameter and supplies implicit order useful for sequential encoders.Subgraph tokenization can avoid a separate ordering step when its order is implicit.
  • Graph Mamba Networks: Bidirectional selective SSMs scan sequences in two directions, improving robustness to token permutations when tokenization lacks implicit order.The model uses Mamba-style selection to control which nodes or subgraphs flow into hidden states.
  • Results: GMNs are theoretically justified as universal graph-function approximators and, with proper PE/SE, more expressive than any WL test.Experiments report strong performance across long-range, small-scale, large-scale, and heterophilic benchmarks with less GPU memory.

2. Related Work and Backgrounds

The background contrasts local message-passing, global-attention graph models, subgraph tokenization, and state space models. It motivates GMNs as an efficient alternative that combines sequence-modeling efficiency with graph-specific processing.

  • Message-Passing Neural Networks: MPNNs iteratively aggregate local neighborhood information but face expressivity, over-smoothing, and over-squashing limitations.Their expressivity is described as limited by the 1-WL isomorphism test.
  • Graph Transformers: Graph Transformers use global attention to connect node pairs, but full attention has quadratic computational cost that limits large-scale applicability.Sparse attention and subgraph tokenization are discussed as efficiency-oriented alternatives.
  • Subgraph Tokenization: Subgraph tokenization represents graphs with extracted local subgraphs, typically initialized through MPNNs, and reduces complexity from node-level attention to O(k^2) for k tokens.These methods can also reduce reliance on complex positional or structural encodings.
  • State Space Models: SSMs map input sequences to response sequences through latent states and state-transition and projection parameters.Discrete SSMs use a parameter ∆ to discretize the continuous system.
  • State Space Models: Discrete-time SSMs can be represented as convolutions and computed efficiently, supporting their use as scalable alternatives to attention.Structured state space models further improve SSM efficiency and scalability through reparameterization.

3. Challenges & Motivations: Transformers vs Mamba

Adapting Mamba to graphs requires addressing the mismatch between recurrent sequence processing and unordered, non-causal graph data. The paper identifies ordering, long-range context, scalability, and tokenization as central design challenges.

  • Sequences vs 2-D Data: Unlike permutation-equivariant Transformers, Mamba scans tokens recurrently and therefore requires an ordering mechanism for graph tokens.Graph adaptations of Transformers can rely on tokenization plus PE/SE, whereas sequential encoders require explicit ordering.
  • Sequences vs 2-D Data: Order sensitivity can also exploit hierarchical neighborhood structure and domain knowledge when token order carries meaning.The paper highlights k-hop neighborhoods as an implicitly ordered example.
  • Long-range Sequence Modeling: Mamba’s selection mechanism can filter irrelevant information and reset its state, enabling longer graph-token sequences to provide additional context.The paper connects longer token sequences with potential gains in context and expressive power.
  • Scalability: GMNs target linear time and memory complexity, making the cost of positional and structural encodings more consequential than in quadratic-cost Transformers.The paper notes that PE/SE construction can become a bottleneck when the encoder itself is linear.
  • Node or Subgraph?: Node tokenization better captures long-range dependencies, whereas subgraph tokenization better captures local neighborhoods and is more efficient in practice.GMNs make this choice tunable through a single parameter m.

4. Graph Mamba Networks

Graph Mamba Networks adapt selective state space models to graphs through neighborhood tokenization, ordered tokens, local encoding, and bidirectional scanning. The framework supports flexible tokenization, theoretical expressiveness guarantees, and long-range information flow.

  • Tokenization and Encoding: GMNs use neighborhood sampling to convert each node’s local and global graph context into a sequence of encoded subgraph tokens.The method samples subgraphs from random walks at multiple lengths and encodes them into node representations.
  • Tokenization and Encoding: GMNs repeat neighborhood sampling s times to create longer, diverse token sequences that exploit selective SSMs’ linear computational cost.The resulting sequence can be longer than the graph diameter and is theoretically more expressive than k-hop neighborhood sampling.
  • Token Ordering: The sampled neighborhoods form an implicitly ordered sequence because longer-hop tokens contain information about shorter-hop neighborhoods.The architecture reverses this order before sequential encoding.
  • Bidirectional Mamba: Bidirectional recurrent scans propagate information across nodes, capture long-range dependencies, and bridge node and subgraph tokenization.The framework can switch between tokenization choices through architectural flexibility, including the m = 0 node-tokenization case.
  • Theoretical Analysis: Positional and structural encodings are optional, while GMNs remain universal approximators and can exceed WL-test expressiveness under stated conditions.With positional encoding, GMNs can distinguish non-isomorphic graphs; without PE and MPNNs, their expressiveness is not bounded by any WL test.

5. Experiments

GMNs are evaluated across diverse graph benchmarks, against message-passing, transformer, and Mamba-based baselines, with experiments covering accuracy, efficiency, memory, and architectural components. They perform strongly across long-range, large-scale, and heterophilic settings while maintaining favorable efficiency and memory usage.

  • Benchmark Performance: GMNs consistently outperform baselines on datasets requiring long-range dependencies.Their long token sequences support long-range learning, while selection filters irrelevant neighborhood information and random-walk sampling captures diverse hierarchical neighborhoods.
  • Benchmark Performance: GMN- achieves competitive performance without a Transformer, MPNN, or PE/SE.This result supports the claim that these components can be sufficient but are not necessary for strong graph-learning performance.
  • Benchmark Performance: GMN and Exphormer achieve competitive performance on GNN Benchmark datasets, with each outperforming the other two times.GMN consistently outperforms the GPS + Mamba baseline, highlighting the importance of a graph-specific framework.
  • Benchmark Performance: GMN outperforms baselines on 3 of 4 heterophilic datasets and achieves the second-best result on the remaining dataset.The authors associate this robustness with selection of relevant information and modeling of long-range dependencies.
  • Ablation Study: All GMN components contribute to performance, with bidirectional Mamba providing the largest contribution in the ablation study.The study removes or modifies bidirectional Mamba, MPNN, token ordering, and PE while holding other components fixed.
  • Efficiency and Memory: GMN variants are the most efficient methods on OGBN-Arxiv and MalNet-Tiny while achieving the best performance.On MalNet-Tiny, GMN memory scales linearly with input size, unlike the higher GPU-memory cost of GPS.

6. Conclusion

The paper introduces Graph Mamba Networks, a graph-learning framework based on state space models, and develops a structured recipe for adapting selective SSMs to graph data. It combines theoretical justification with empirical evaluation of the proposed design.

  • Conclusion: Graph Mamba Networks are introduced as a new class of graph-learning models based on state space models.The paper studies challenges in adapting SSMs to graph-structured data and evaluates GMNs experimentally.
  • Conclusion: The GMN design uses four required steps and one optional step: neighborhood tokenization, token ordering, bidirectional selective SSM encoding, local encoding, and PE/SE.PE and SE are described as dispensable, while the other four components form the required recipe.

A. Details of Datasets

The paper reports dataset statistics separately and directs readers to the cited benchmark sources for additional dataset details.

  • Dataset Details: Dataset statistics are reported in Table 5, with further details referenced from four established graph-learning benchmarks.The cited sources cover long-range, GNN, heterophilic, and Open Graph Benchmark datasets.

B. Experimental Setup

The experiments use grid-searched hyperparameters, established data splits, multiple random seeds, and reported results from prior benchmark studies for baselines.

  • Experimental Setup: Hyperparameters are tuned by grid search, using the search space listed in Table 6.The experiments follow prior studies for training, validation, and test splits.
  • Experimental Setup: Results are reported over 4 random seeds using the same training, validation, and test splits as prior work.Baseline results in Tables 1–3 are reused from several earlier studies.

C. Details of GMN Architecture: Algorithms

The GMN algorithm constructs graph tokens from sampled walks and uses configurable hyperparameters, with some preparation performed before training. The reported setup also reuses prior benchmark results and notes an incompletely searched space.

  • GMNs sample walks of lengths 1 through m for each node, construct induced-subgraph tokens, and repeat the process s times for longer sequences and more samples.
  • The hyperparameter search space is reported in Table 6, while the supplied results note that this space was not fully searched.

D. Additional Experimental Results

Additional experiments examine how GMN tokenization parameters affect performance and compare GMNs with neighborhood-aggregation baselines. Performance depends on token construction and sequence length, while Mamba-based selection and sampled walks provide the stated comparison rationale.

  • Effect of M: Performance peaks at a dataset-dependent value of M, the number of walks aggregated to construct each subgraph token.The experiment varies M from 1 to 10 on Roman-empire and PascalVOC-SP.
  • Effect of m: Increasing m does not damage performance and may improve it when tokenization includes farther neighbors.The reported experiments vary m from 1 to 60 across Roman-empire and PascalVOC-SP.
  • Effect of s: Increasing s monotonically improves performance because longer token sequences provide more context while Mamba selects informative tokens and filters irrelevant ones.The effect is reported in Figure 4 (Right).
  • Comparison with related methods: GMNs use sampled walks and Mamba’s selection mechanism, whereas GRED and S4G aggregate distance-based neighborhoods with an RNN or structured SSM.The comparison concerns GMNs, S4G, and GRED on the datasets used by the latter methods.

E. Complexity Analysis of GNMs

GMN complexity is linear in graph size for fixed tokenization parameters, but longer token sequences create a performance–runtime trade-off. Bidirectional Mamba contributes linear sequence processing, while optional MPNN layers preserve linear graph-time complexity.

  • Complexity: For m ≥ 1, GMN time complexity is O(M × s × (m + 1) × |V| + |E|), linear in |V| and |E|.The per-node walk generation and bidirectional Mamba processing use O(M × s × (m + 1)) time.
  • Complexity: Bidirectional Mamba processes the token sequence with twice the complexity of Mamba while remaining linear in the number of tokens.The algorithm takes an ordered token matrix as input and returns updated sequence encodings.
  • Complexity: Optional MPNN integration does not change the asymptotic complexity because MPNNs require O(|V| + |E|) time.This statement applies when the MPNN is used as an optional architecture component.
  • Complexity: When m = 0, each node is a token and GMN requires O(|V|) time, or O(|V| + |E|) with an MPNN.In this case, MPNN time dominates the architecture’s complexity.
  • Trade-off: Larger s can improve performance but also makes the model slower, creating a trade-off between accuracy and time complexity.The stated relationship applies for m ≥ 1.

F. Discussion on a Concurrent Work

The discussion contrasts GMNs with a concurrent Mamba-based GPS variant, emphasizing differences in tokenization, directionality, positional encodings, and expressive power. GMNs retain flexible tokenization and bidirectional processing, while the comparison identifies limitations of the concurrent design.

  • Architecture Design: The concurrent GMB method replaces GPS’s attention module with Mamba, treats each node as a token, uses PE/SE, and orders nodes by degree.The paper states that this design inherits limitations associated with node tokenization and one-directional Mamba.
  • Architecture Design: GMNs generalize GMB by allowing node or subgraph tokenization through m and by using bidirectional rather than one-directional Mamba.GMB corresponds to the special case m = 0 with one-directional Mamba.
  • Architecture Design: GMNs’ parameter m provides flexibility to prioritize inductive bias or long-range dependencies according to the task and dataset.The paper presents this flexibility as a distinction from GMB.
  • Expressive Power: GMB requires complex PE/SE because it lacks inductive bias, whereas GMNs without PE/SE and MPNNs are stated to have unbounded expressive power with respect to the isomorphic test.This comparison is presented as part of the expressive-power discussion.
  • Expressive Power: GMB cannot distinguish graphs with the same degree sequence when node features are also identical, so its expressive power is bounded by its MPNN.The paper gives this as a worst-case limitation of the GMB input representation.

G. Theoretical Analysis of GMNs

Theoretical results establish that GMNs are universal approximators and can exceed the expressive power of bounded WL tests, both with and, under stated conditions, without positional encodings. The analysis also shows that their neighborhood sampling can be strictly more expressive than k-hop sampling.

  • Neighborhood sampling: GMNs’ neighborhood sampling is strictly more expressive than k-hop neighborhood sampling when M, m, and s are sufficiently large.The result is stated as Theorem 4.1, with the appendix proof relating random-walk sampling to k-hop sampling.
  • Universality: For any permutation-equivariant continuous function, a GMN with positional encoding can approximate it arbitrarily closely in the stated ℓp sense.The universality theorem quantifies approximation by requiring ℓp(f, g) < ϵ for any ϵ > 0.
  • Expressive power with PE: With the full set of eigenfunctions and enough parameters, GMNs can distinguish every pair of non-isomorphic graphs and are more powerful than any WL test.This expressive-power claim is stated for GMNs using positional encoding.
  • Expressive power without PE and MPNN: Without positional encoding or MPNNs, GMNs remain unbounded in expressive power relative to WL tests: for every k, some graphs separate under GMNs but not k-WL.The proof connects GMNs using random-walk features to CRaWl and its one-dimensional convolutional processing.
Loading 2402.08678v2…