Graph Mamba: В направлении обучения на графах с использованием моделей пространства состояний

Graph Mamba: Towards Learning on Graphs with State Space Models

February 13, 2024

Авторы: Ali Behrouz, Farnoosh Hashemi

cs.AI

Аннотация

Графовые нейронные сети (GNNs) продемонстрировали значительный потенциал в обучении представлений графов. Большинство GNNs определяют локальный механизм передачи сообщений, распространяя информацию по графу с помощью наложения нескольких слоев. Однако эти методы известны двумя основными ограничениями: чрезмерным сжатием информации и слабым учетом дальнодействующих зависимостей. Недавно Графовые Трансформеры (GTs) появились как мощная альтернатива нейронным сетям с передачей сообщений (MPNNs). Тем не менее, GTs имеют квадратичную вычислительную сложность, не обладают индуктивными смещениями на структуры графов и полагаются на сложные позиционные/структурные кодирования (SE/PE). В данной работе мы показываем, что хотя Трансформеры, сложная передача сообщений и SE/PE достаточны для хорошей производительности на практике, ни один из этих элементов не является необходимым. Вдохновленные недавним успехом моделей пространства состояний (SSMs), таких как Mamba, мы представляем Графовые сети Mamba (GMNs), общую основу для нового класса GNNs, основанных на селективных SSMs. Мы обсуждаем и классифицируем новые вызовы при адаптации SSMs к графовым данным и представляем четыре обязательных и один необязательный шаг для проектирования GMNs, где мы выбираем (1) Токенизацию окрестностей, (2) Упорядочивание токенов, (3) Архитектуру двунаправленного селективного SSM-кодировщика, (4) Локальное кодирование и необязательное (5) PE и SE. Мы также предоставляем теоретическое обоснование мощности GMNs. Эксперименты показывают, что, несмотря на значительно меньшую вычислительную стоимость, GMNs достигают выдающейся производительности на наборах данных с дальнодействующими, малыми, крупными и гетерофильными характеристиками.

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.