Graph Mamba: Rumo à Aprendizagem em Grafos com Modelos de Espaço de Estados

Graph Mamba: Towards Learning on Graphs with State Space Models

February 13, 2024

Autores: Ali Behrouz, Farnoosh Hashemi

cs.AI

Resumo

Redes Neurais em Grafos (GNNs) têm demonstrado um potencial promissor no aprendizado de representações em grafos. A maioria das GNNs define um mecanismo local de passagem de mensagens, propagando informações pelo grafo por meio da empilhamento de múltiplas camadas. Esses métodos, no entanto, são conhecidos por sofrer de duas grandes limitações: sobrecompressão e dificuldade em capturar dependências de longo alcance. Recentemente, Transformers em Grafos (GTs) surgiram como uma alternativa poderosa às Redes Neurais de Passagem de Mensagens (MPNNs). GTs, porém, têm um custo computacional quadrático, carecem de vieses indutivos sobre estruturas de grafos e dependem de Codificações Posicionais/Estruturais (SE/PE) complexas. Neste artigo, mostramos que, embora Transformers, passagem de mensagens complexa e SE/PE sejam suficientes para um bom desempenho na prática, nenhum deles é necessário. Motivados pelo recente sucesso dos Modelos de Espaço de Estados (SSMs), como o Mamba, apresentamos as Redes Mamba em Grafos (GMNs), uma estrutura geral para uma nova classe de GNNs baseada em SSMs seletivos. Discutimos e categorizamos os novos desafios ao adotar SSMs para dados estruturados em grafos e apresentamos quatro etapas necessárias e uma opcional para projetar GMNs, onde escolhemos (1) Tokenização de Vizinhança, (2) Ordenação de Tokens, (3) Arquitetura do Codificador Bidirecional de SSM Seletivo, (4) Codificação Local, e dispensável (5) PE e SE. Além disso, fornecemos justificativas teóricas para o poder das GMNs. Experimentos demonstram que, apesar de um custo computacional muito menor, as GMNs alcançam um desempenho excepcional em conjuntos de dados de referência de longo alcance, pequena escala, grande escala e heterofílicos.

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.