Graph Mamba: Verso l'Apprendimento su Grafi con Modelli a Spazio di Stato

Graph Mamba: Towards Learning on Graphs with State Space Models

February 13, 2024

Autori: Ali Behrouz, Farnoosh Hashemi

cs.AI

Abstract

Le Graph Neural Networks (GNN) hanno dimostrato un potenziale promettente nell'apprendimento di rappresentazioni di grafi. La maggior parte delle GNN definisce un meccanismo locale di passaggio di messaggi, propagando le informazioni sul grafo attraverso l'impilamento di più strati. Tuttavia, questi metodi sono noti per soffrire di due principali limitazioni: l'eccessiva compressione (over-squashing) e la scarsa capacità di catturare dipendenze a lungo raggio. Recentemente, i Graph Transformers (GT) sono emersi come una potente alternativa alle Message-Passing Neural Networks (MPNN). I GT, tuttavia, presentano un costo computazionale quadratico, mancano di bias induttivi sulle strutture dei grafi e si basano su complessi Positional/Structural Encodings (SE/PE). In questo articolo, dimostriamo che, sebbene i Transformers, il complesso passaggio di messaggi e gli SE/PE siano sufficienti per ottenere buone prestazioni nella pratica, nessuno di questi elementi è strettamente necessario. Motivati dal recente successo dei State Space Models (SSM), come Mamba, presentiamo le Graph Mamba Networks (GMN), un framework generale per una nuova classe di GNN basata su SSM selettivi. Discutiamo e categorizziamo le nuove sfide nell'adottare gli SSM per dati strutturati a grafo e presentiamo quattro passaggi obbligatori e uno opzionale per progettare le GMN, in cui scegliamo (1) Tokenizzazione del Vicinato, (2) Ordinamento dei Token, (3) Architettura del Codificatore SSM Selettivo Bidirezionale, (4) Codifica Locale, e opzionalmente (5) PE e SE. Forniamo inoltre una giustificazione teorica per la potenza delle GMN. Gli esperimenti dimostrano che, nonostante un costo computazionale molto inferiore, le GMN raggiungono prestazioni eccezionali su dataset di benchmark a lungo raggio, su piccola scala, su larga scala e eterofili.

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.