Graph Mamba: Naar Leren op Grafieken met State Space Modellen

Graph Mamba: Towards Learning on Graphs with State Space Models

February 13, 2024

Auteurs: Ali Behrouz, Farnoosh Hashemi

cs.AI

Samenvatting

Graph Neural Networks (GNNs) hebben veelbelovend potentieel getoond in het leren van grafische representaties. De meeste GNNs definiëren een lokaal berichtenoverdrachtsmechanisme, waarbij informatie over de grafiek wordt verspreid door meerdere lagen te stapelen. Deze methoden staan echter bekend om twee belangrijke beperkingen: over-squashing en een slechte vastlegging van afhankelijkheden over lange afstanden. Recentelijk zijn Graph Transformers (GTs) naar voren gekomen als een krachtig alternatief voor Message-Passing Neural Networks (MPNNs). GTs hebben echter een kwadratische rekencapaciteit, missen inductieve vooroordelen op grafiekstructuren en zijn afhankelijk van complexe Positionele/Structurele Coderingen (SE/PE). In dit artikel tonen we aan dat hoewel Transformers, complexe berichtenoverdracht en SE/PE in de praktijk voldoende zijn voor goede prestaties, geen van deze noodzakelijk is. Gemotiveerd door het recente succes van State Space Models (SSMs), zoals Mamba, presenteren we Graph Mamba Networks (GMNs), een algemeen raamwerk voor een nieuwe klasse van GNNs gebaseerd op selectieve SSMs. We bespreken en categoriseren de nieuwe uitdagingen bij het toepassen van SSMs op grafisch gestructureerde gegevens en presenteren vier vereiste en één optionele stap om GMNs te ontwerpen, waarbij we kiezen voor (1) Neighborhood Tokenization, (2) Token Ordering, (3) Architectuur van Bidirectionele Selectieve SSM Encoder, (4) Lokale Codering, en optioneel (5) PE en SE. We bieden verder een theoretische rechtvaardiging voor de kracht van GMNs. Experimenten tonen aan dat GMNs, ondanks veel minder rekencapaciteit, een uitstekende prestatie leveren in benchmarksets voor lange afstanden, kleine schaal, grote schaal en heterofiele gegevens.

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.