Graph Mamba: Auf dem Weg zum Lernen auf Graphen mit State-Space-Modellen

Graph Mamba: Towards Learning on Graphs with State Space Models

February 13, 2024

Autoren: Ali Behrouz, Farnoosh Hashemi

cs.AI

Zusammenfassung

Graph Neural Networks (GNNs) haben vielversprechendes Potenzial im Bereich des Graph-Representation-Learnings gezeigt. Die Mehrheit der GNNs definiert einen lokalen Nachrichtenübermittlungsmechanismus, der Informationen über den Graphen durch das Stapeln mehrerer Schichten weiterleitet. Diese Methoden leiden jedoch bekanntermaßen unter zwei wesentlichen Einschränkungen: Over-Squashing und einer unzureichenden Erfassung von langreichweitigen Abhängigkeiten. Kürzlich sind Graph Transformers (GTs) als leistungsstarke Alternative zu Message-Passing Neural Networks (MPNNs) aufgetaucht. GTs haben jedoch quadratische Rechenkosten, fehlende induktive Verzerrungen auf Graphstrukturen und sind auf komplexe Positional/Structural Encodings (SE/PE) angewiesen. In diesem Artikel zeigen wir, dass zwar Transformers, komplexe Nachrichtenübermittlung und SE/PE in der Praxis für eine gute Leistung ausreichen, jedoch keines davon notwendig ist. Motiviert durch den jüngsten Erfolg von State Space Models (SSMs), wie Mamba, präsentieren wir Graph Mamba Networks (GMNs), ein allgemeines Framework für eine neue Klasse von GNNs, die auf selektiven SSMs basieren. Wir diskutieren und kategorisieren die neuen Herausforderungen bei der Anwendung von SSMs auf graphstrukturierte Daten und stellen vier erforderliche und einen optionalen Schritt zur Gestaltung von GMNs vor, bei denen wir (1) Neighborhood Tokenization, (2) Token Ordering, (3) Architektur des Bidirektionalen Selektiven SSM Encoders, (4) Local Encoding und optional (5) PE und SE wählen. Wir liefern weiterhin eine theoretische Begründung für die Leistungsfähigkeit von GMNs. Experimente zeigen, dass GMNs trotz deutlich geringerer Rechenkosten eine herausragende Leistung in langreichweitigen, kleinskaligen, großskaligen und heterophilen Benchmark-Datensätzen erzielen.

English

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 adopting 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.