Graph Mamba : Vers l'apprentissage sur les graphes avec des modèles à espace d'états

Graph Mamba: Towards Learning on Graphs with State Space Models

February 13, 2024

Auteurs: Ali Behrouz, Farnoosh Hashemi

cs.AI

Résumé

Les réseaux de neurones sur graphes (Graph Neural Networks, GNNs) ont montré un potentiel prometteur dans l'apprentissage de représentations de graphes. La majorité des GNNs définissent un mécanisme local de passage de messages, propageant l'information sur le graphe en empilant plusieurs couches. Ces méthodes, cependant, sont connues pour souffrir de deux limitations majeures : le sur-écrasement (over-squashing) et une mauvaise capture des dépendances à longue portée. Récemment, les Transformers sur graphes (Graph Transformers, GTs) sont apparus comme une alternative puissante aux réseaux de neurones à passage de messages (Message-Passing Neural Networks, MPNNs). Les GTs, cependant, ont un coût computationnel quadratique, manquent de biais inductifs sur les structures de graphes, et reposent sur des encodages positionnels/structurels complexes (Positional/Structural Encodings, SE/PE). Dans cet article, nous montrons que bien que les Transformers, le passage de messages complexe et les SE/PE soient suffisants pour une bonne performance en pratique, aucun n'est nécessaire. Motivés par le succès récent des modèles d'espace d'états (State Space Models, SSMs), tels que Mamba, nous présentons les réseaux Graph Mamba (Graph Mamba Networks, GMNs), un cadre général pour une nouvelle classe de GNNs basée sur des SSMs sélectifs. Nous discutons et catégorisons les nouveaux défis lors de l'adoption des SSMs pour les données structurées en graphes, et présentons quatre étapes nécessaires et une étape optionnelle pour concevoir les GMNs, où nous choisissons (1) la tokenisation du voisinage, (2) l'ordonnancement des tokens, (3) l'architecture de l'encodeur bidirectionnel SSM sélectif, (4) l'encodage local, et l'étape dispensable (5) PE et SE. Nous fournissons en outre une justification théorique de la puissance des GMNs. Les expériences démontrent que malgré un coût computationnel bien moindre, les GMNs atteignent une performance exceptionnelle sur des ensembles de données de référence à longue portée, à petite échelle, à grande échelle et hétérophiles.

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.