Graph Mamba: Hacia el Aprendizaje en Grafos con Modelos de Espacio de Estados

Graph Mamba: Towards Learning on Graphs with State Space Models

February 13, 2024

Autores: Ali Behrouz, Farnoosh Hashemi

cs.AI

Resumen

Las Redes Neuronales de Grafos (GNNs) han demostrado un potencial prometedor en el aprendizaje de representaciones de grafos. La mayoría de las GNNs definen un mecanismo local de paso de mensajes, propagando información a través del grafo mediante la superposición de múltiples capas. Sin embargo, se sabe que estos métodos sufren dos limitaciones principales: el sobre-aplanamiento y la captura deficiente de dependencias de largo alcance. Recientemente, los Transformers de Grafos (GTs) surgieron como una alternativa poderosa a las Redes Neuronales de Paso de Mensajes (MPNNs). No obstante, los GTs tienen un costo computacional cuadrático, carecen de sesgos inductivos sobre las estructuras de grafos y dependen de Codificaciones Posicionales/Estructurales (SE/PE) complejas. En este artículo, demostramos que, aunque los Transformers, el paso de mensajes complejo y las SE/PE son suficientes para un buen rendimiento en la práctica, ninguno de ellos es necesario. Motivados por el reciente éxito de los Modelos de Espacio de Estados (SSMs), como Mamba, presentamos las Redes Mamba de Grafos (GMNs), un marco general para una nueva clase de GNNs basadas en SSMs selectivos. Discutimos y categorizamos los nuevos desafíos al adoptar SSMs para datos estructurados en grafos, y presentamos cuatro pasos requeridos y uno opcional para diseñar GMNs, donde elegimos (1) Tokenización del Vecindario, (2) Ordenamiento de Tokens, (3) Arquitectura del Codificador Bidireccional de SSM Selectivo, (4) Codificación Local, y el dispensable (5) PE y SE. Además, proporcionamos una justificación teórica sobre el poder de las GMNs. Los experimentos demuestran que, a pesar de un costo computacional mucho menor, las GMNs alcanzan un rendimiento sobresaliente en conjuntos de datos de referencia de largo alcance, pequeña escala, gran escala y heterofilia.

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.