그래프 맘바: 상태 공간 모델을 활용한 그래프 학습을 향하여

초록

그래프 신경망(GNNs)은 그래프 표현 학습에서 유망한 잠재력을 보여주고 있다. 대부분의 GNN은 로컬 메시지 전달 메커니즘을 정의하며, 여러 층을 쌓아 그래프 위에서 정보를 전파한다. 그러나 이러한 방법들은 두 가지 주요 한계로 인해 어려움을 겪는 것으로 알려져 있다: 과도한 정보 압축(over-squashing)과 장거리 의존성을 효과적으로 포착하지 못하는 문제이다. 최근, 그래프 트랜스포머(GTs)가 메시지 전달 신경망(MPNNs)의 강력한 대안으로 등장했다. 그러나 GTs는 이차 계산 비용이 발생하며, 그래프 구조에 대한 귀납적 편향(inductive biases)이 부족하고, 복잡한 위치/구조 인코딩(SE/PE)에 의존한다. 본 논문에서는 트랜스포머, 복잡한 메시지 전달, 그리고 SE/PE가 실제로 좋은 성능을 내기에 충분하지만, 이들 중 어느 것도 필수적이지 않음을 보인다. 최근 Mamba와 같은 상태 공간 모델(SSMs)의 성공에 영감을 받아, 우리는 선택적 SSMs를 기반으로 한 새로운 클래스의 GNN인 그래프 맘바 네트워크(GMNs)를 제안한다. 우리는 그래프 구조 데이터에 SSMs를 적용할 때 발생하는 새로운 도전 과제들을 논의하고 분류하며, GMNs를 설계하기 위해 필요한 네 가지 단계와 선택적인 한 가지 단계를 제시한다. 이 단계들은 (1) 이웃 토큰화, (2) 토큰 순서화, (3) 양방향 선택적 SSM 인코더의 아키텍처, (4) 로컬 인코딩, 그리고 선택적인 (5) PE와 SE를 포함한다. 또한, 우리는 GMNs의 강력함에 대한 이론적 근거를 제공한다. 실험 결과, GMNs는 훨씬 적은 계산 비용에도 불구하고 장거리, 소규모, 대규모, 그리고 이질적(heterophilic) 벤치마크 데이터셋에서 뛰어난 성능을 달성함을 보여준다.

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.

그래프 맘바: 상태 공간 모델을 활용한 그래프 학습을 향하여

Graph Mamba: Towards Learning on Graphs with State Space Models

초록

Support