Graph Mamba: Towards Learning on Graphs with State Space Models

  • 2024-02-19 18:53:13
  • Ali Behrouz, Farnoosh Hashemi
  • 0

Abstract

Graph Neural Networks (GNNs) have shown promising potential in graphrepresentation learning. The majority of GNNs define a local message-passingmechanism, 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, GraphTransformers (GTs) emerged as a powerful alternative to Message-Passing NeuralNetworks (MPNNs). GTs, however, have quadratic computational cost, lackinductive biases on graph structures, and rely on complex Positional/StructuralEncodings (SE/PE). In this paper, we show that while Transformers, complexmessage-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 generalframework for a new class of GNNs based on selective SSMs. We discuss andcategorize the new challenges when adapting SSMs to graph-structured data, andpresent four required and one optional steps to design GMNs, where we choose(1) Neighborhood Tokenization, (2) Token Ordering, (3) Architecture ofBidirectional Selective SSM Encoder, (4) Local Encoding, and dispensable (5) PEand SE. We further provide theoretical justification for the power of GMNs.Experiments demonstrate that despite much less computational cost, GMNs attainan outstanding performance in long-range, small-scale, large-scale, andheterophilic benchmark datasets.

 

Quick Read (beta)

loading the full paper ...