PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chloe Wang, Oleksii Tsepa, Jun Ma, Bo Wang

arXiv:2402.00789v1cs.LGcs.AI

TL;DR

Long-range graph Transformers face quadratic attention costs, and existing sparsification methods provide limited data-dependent context reasoning. Graph-Mamba integrates selective state spaces with graph-specific node prioritization and permutation, achieving strong performance across ten datasets with reduced computational and memory costs.

  • Problem

    Graph Transformers model long-range dependencies but scale poorly because attention is quadratic, while existing graph sparsification lacks data-dependent context reasoning.

  • Method

    Graph-Mamba replaces GraphGPS attention with a selective state-space block using input-dependent node filtering, node prioritization, and permutation-based graph adaptation.

  • Results

    Graph-Mamba achieves state-of-the-art or comparable performance across ten datasets while using linear-time computation and a fraction of the computational cost.

  • Takeaways & Limitations

    Graph-Mamba offers a promising alternative to dense or sparse graph attention for long-range prediction with competitive predictive power and context awareness.

  • Takeaways & Limitations

    Effective ways to inject graph topology into flattened input sequences remain unexplored, and learning the optimal flattening strategy from data is left for future work.

Abstract

from arXiv · show

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.

1. Introduction

Graph Transformers capture long-range graph dependencies but face quadratic attention costs, while existing sparsification methods lack data-dependent context reasoning. Graph-Mamba addresses this by integrating selective state spaces with graph-specific node prioritization and permutation, achieving improved performance and efficiency.

  • Graph Transformer attention lets each node interact with all others, but its quadratic cost hinders scaling to long sequences and large graphs.
  • BigBird’s sequence-oriented sparsification does not generalize well to graphs, causing performance deterioration in GraphGPS.
  • Mamba selects which input context enters hidden states, providing a data-dependent node-selection view for graph modeling.
  • Graph-Mamba integrates a selective state space model as a data-dependent alternative to Graph Transformer attention sparsification.
  • Node prioritization gives important nodes more context, while permutation reduces sequence-related biases when adapting state spaces to non-sequential graphs.
  • Graph-Mamba achieves linear-time complexity and reduces GPU memory consumption by up to 74% on large graphs across ten public datasets.

2. Related Work

Graph modeling methods range from local message passing to dense and sparse attention, while state space models offer recurrent, efficient sequence processing. Their differing strengths motivate graph-specific approaches to long-range dependency modeling.

  • GNNs use message passing to aggregate information from neighboring nodes, but MPNNs have expressive power limited by the 1-dimensional Weisfeiler-Lehman framework.
  • Graph Transformers use full attention to capture long-range dependencies beyond local neighborhoods, at O(N^2) computational complexity.
  • GraphGPS combines structural and positional encodings, an MPNN, and a graph transformer in a modular framework.
  • BigBird and Performer reduce attention costs through graph or token subsampling and lower-dimensional kernel approximations, respectively.
  • State space models recurrently store sequence context in hidden states and update outputs using those states together with current inputs.

3. Graph-Mamba

Graph-Mamba replaces GraphGPS attention with a selective state-space block that performs input-dependent context filtering and graph-specific node ordering. Its recurrent design targets efficient long-range modeling while combining local message passing with global node updates.

  • Graph-Mamba framework: Graph-Mamba inserts a Graph-Mamba block into GraphGPS, using recurrent scanning and selection for input-dependent graph sparsification.
  • Graph-specific sparsification: The selective mechanism filters relevant long-range information, while node prioritization gives important nodes greater access to prior context.
  • Structured state space models: An SSM maps an input sequence to an output sequence through a latent hidden state, with A, B, and C as state, input, and output matrices.
  • Structured state space models: Discretization produces state parameters Ā and B̄ from continuous dynamics using step size ∆, enabling computation on discrete real-world data.
  • Selective state spaces: Mamba makes SSM parameters input-dependent, allowing current nodes to control hidden-state updates and filter irrelevant prior context.
  • Node prioritization: Nodes are heuristically sorted before selective SSM processing and restored to their original order afterward.
  • Graph-Mamba block: The GMB layer combines edge-based GatedGCN message passing with node-focused GMB updates, producing updated node and edge embeddings.
  • Evaluation: Table 1 benchmarks Graph-Mamba against existing methods on five long-range datasets containing 150 to 1,400 nodes per graph.

4. Experiments

Experiments across ten graph datasets show that Graph-Mamba maintains strong predictive performance while reducing computational and memory costs. Its permutation and node-prioritization recipe improves performance, especially on long-range graph tasks.

  • Benchmark setup: Graph-Mamba was evaluated on ten datasets covering graph-, node-, and link-level classification and regression tasks.Test metrics were reported across multiple runs.
  • Predictive performance: In four of five datasets with 150–1,400 nodes per graph, Graph-Mamba improved performance over sparse attention methods by up to 5%.It also performed comparably to dense Transformer and sparse attention methods on smaller graphs.
  • Efficiency: Graph-Mamba scales linearly in FLOPs and memory with input length, whereas GPS-Transformer scales quadratically.On MalNet-Tiny, Graph-Mamba trained full graphs with twice the number of nodes at batch sizes up to 256, while GPS-Transformer encountered out-of-memory issues below 700 nodes at batch size 16.
  • Efficiency: On Peptides-func, Graph-Mamba reduced memory usage by 74% and FLOPs by 66% versus Transformer.It also reduced both FLOPs and memory usage by 40% versus Exphormer.
  • Ablation: Node-level permutation increased average precision by 3% on Peptides-Func and F1 scores by 10% on PascalVOC-SP versus the baseline.Combining degree-based node prioritization with permutation further improved PascalVOC-SP scores.

5. Conclusion

Graph-Mamba uses selective state spaces for data-dependent context selection in graphs, together with graph-specific permutation and node-prioritization strategies. Across ten datasets, it achieved state-of-the-art or comparable performance with linear-time computation and reduced memory use, while several sequence-construction questions remain open.

  • Method: Graph-Mamba uses an SSM selection mechanism to filter relevant nodes, compress long-range context, and propagate it during node-embedding updates.A recurrent scan provides linear-time complexity and reduced memory consumption.
  • Results: Across ten datasets and varied graph prediction tasks, Graph-Mamba achieved state-of-the-art or comparable performance at a fraction of dense or sparse attention’s computational cost.The reported conclusion also highlights long-range context awareness.
  • Limitations and future work: Future work should explore architectures beyond GraphGPS and learn more effective ways to flatten graphs into sequences and inject topology.The paper identifies optimal data-driven sequence construction as unresolved.

Impact Statement

The paper identifies potential adverse or malicious applications of Graph-Mamba in domains such as drug discovery and healthcare. It does not attribute direct negative societal outcomes to the general representation-learning method itself.

  • Societal impact: Although Graph-Mamba is designed as a general graph representation-learning method, adverse or malicious applications in drug discovery and healthcare may cause undesirable effects.The impact statement distinguishes the method’s general design from risks arising in particular applications.

A. Dataset Description.

The evaluation spans ten graph datasets covering image, software, synthetic community, molecular, and superpixel-based tasks across graph, node, and link prediction settings.

  • Benchmark scope: The benchmark covers ten datasets from LRGB and GNN Benchmark, with five long-input datasets and five small-to-medium-input datasets.The long-input group corresponds to Table 1, while the smaller group corresponds to Table 4.
  • Image graphs: CIFAR10 and MNIST represent images as 8-nearest-neighbor graphs of SLIC superpixels with ten classification labels.Both use the standard splits of their original image datasets.
  • Software graphs: MalNet-Tiny contains 5,000 Android function-call graphs with up to 5,000 nodes, classified into one benign and four malware categories using local degree profiles.Original node and edge features are removed, making the task rely solely on graph structure.
  • Synthetic graphs: PATTERN and CLUSTER are inductive node-classification datasets generated from stochastic block models with subgraph-pattern and cluster-identity objectives.PATTERN distinguishes 100 subgraph patterns, whereas CLUSTER predicts six cluster identities from designated test nodes.
  • Molecular graphs: Peptides-Func and Peptides-Struct use peptide graphs with large diameters for graph-level functional classification and structural-property regression.Peptides-Func has ten functional labels, while Peptides-Struct predicts eleven molecular structural properties.
  • Molecular graphs: PCQM-Contact evaluates link prediction for distant molecular nodes that are spatially close in three dimensions using Mean Reciprocal Rank.The task is based on PCQM4Mv2 molecular structures.
  • Superpixel graphs: PascalVOC-SP and COCO-SP classify SLIC superpixels into object classes through node-level graph prediction tasks analogous to semantic segmentation.Both datasets are graph representations of image datasets.

B. Additional Benchmark Results.

On short- to medium-range graphs, Graph-Mamba maintains predictive performance comparable to GraphGPS with full Transformer attention and Exphormer, supporting generalization to common graph tasks.

  • Evaluation setting: 30–120 nodes per graph characterize the five short- to medium-range datasets evaluated in Table 4.These datasets complement the long-input benchmarks reported separately.
  • Benchmark result: Graph-Mamba achieves predictive performance comparable to GraphGPS with full Transformer attention and Exphormer on these common graph tasks.The result is presented as evidence of generalizability beyond long-range graph settings.

C. Proof of Theorem.

The theorem proof revisits selective SSM discretization under a scalar state and substitutes the input-dependent step size into the zero-order hold formulas.

  • Theorem setup: The theorem considers N = 1, A = -1, B = 1, and ∆t = softplus(Linear(x_t)) for a selective SSM.These choices define the scalar setting used for the recurrence analysis.
  • Discretization: The proof substitutes ∆t into the discretized state-transition formulas for Ā and B̄ under zero-order hold discretization.The formulas are Ā = exp(∆tA) and B̄ = (∆tA)^−1(exp(∆t · A) − I) · ∆tB.

D. Input Node Prioritization and Permutation Strategies.

Graph-Mamba addresses the directional context imbalance of recurrent scans by prioritizing nodes, randomizing ties through permutation, and using variants or binning for different graph structures and sizes.

  • Motivation: A unidirectional Mamba scan gives later nodes access to more accumulated context than earlier nodes in a randomly ordered graph sequence.For example, ND1 updates using hidden states encoding ND0, while later positions incorporate more preceding nodes.
  • Alternative strategies: Eigenvector centrality and cluster-level permutation were explored as alternatives to degree-based prioritization and node-level permutation.Cluster permutation groups nodes by edge-connectivity clusters before randomly permuting nodes within each cluster.
  • Ablation findings: Cluster-level permutation improves over the baseline but performs slightly below node-level permutation, while eigenvector-centrality prioritization is less effective than degree-based prioritization.These comparisons are reported on Peptides-Func and PascalVOC-SP.
  • Large-graph handling: For large graphs, binning divides long node sequences into n sub-sequences, applies prioritization and permutation within each, then recombines the outputs.The technique is used to further sparsify embedding-update connections.
  • Implementation: Graph-Mamba uses matched parameter counts against GPS+Transformer, with a Mamba state dimension of 16, convolution kernel size of 4, and expansion factor of 1.The binning technique is applied only to large datasets with close to or more than 500 nodes.

G. Benchmarking Experiments.

The benchmarking experiments use standardized evaluation pipelines, repeated runs or consolidated benchmark scores, and separate per-example FLOPs and memory estimates. Graph-Mamba is implemented in PyTorch and trained on a single RTX6000 or A100 GPU.

  • Evaluation protocol: Graph-Mamba, GPS+Transformer, and Exphormer benchmark results were averaged over five random seeds, while some earlier-method scores were consolidated from the Exphormer benchmark.The reported metrics followed the same evaluation pipelines as GraphGPS and Exphormer.
  • Evaluation protocol: The benchmark capped maximum GPU memory at 24GB, with an OOM case recorded in the evaluation protocol.
  • Efficiency profiling: FLOPs per sample were estimated from total forward-pass FLOPs in one training epoch divided by the number of training examples.These profiling statistics were collected from a single training epoch.
  • Efficiency profiling: Average memory per training example was estimated by dividing peak GPU memory usage by batch size.The estimate was collected during the FLOPs and memory benchmark.
  • Implementation: Graph-Mamba was implemented with PyTorch, mamba-ssm, and PyTorch Geometric, and trained on a single RTX6000 or A100 GPU.
Loading 2402.00789v1…