[논문 리뷰] Graph-Mamba: Towards Long-Range Graph Sequence Modeling with Selective State Spaces

Xiaoming Wang, Oleksii Tsepa외 2명|arXiv (Cornell University)|2024. 02. 01.

Graph Theory and Algorithms인용 수 20

한 줄 요약

Graph-Mamba는 GraphGPS의 주의 메커니즘을 대체하는 선택적 상태 공간 모델(SSM) 기반의 그래프 블록을 도입하여 데이터 의존적 장거리 컨텍스트 모델링을 선형 시간 복잡도로 가능하게 하고 메모리를 감소시킨다. 이는 밀집 및 희소 주의 방법과 비교해 장거리 그래프 벤치마크에서 우수한 성능을 보인다.

ABSTRACT

Attention mechanisms have been widely used to capture long-range dependencies among nodes in Graph Transformers. Bottlenecked by the quadratic computational cost, attention mechanisms fail to scale in large graphs. Recent improvements in computational efficiency are mainly achieved by attention sparsification with random or heuristic-based graph subsampling, which falls short in data-dependent context reasoning. State space models (SSMs), such as Mamba, have gained prominence for their effectiveness and efficiency in modeling long-range dependencies in sequential data. However, adapting SSMs to non-sequential graph data presents a notable challenge. In this work, we introduce Graph-Mamba, the first attempt to enhance long-range context modeling in graph networks by integrating a Mamba block with the input-dependent node selection mechanism. Specifically, we formulate graph-centric node prioritization and permutation strategies to enhance context-aware reasoning, leading to a substantial improvement in predictive performance. Extensive experiments on ten benchmark datasets demonstrate that Graph-Mamba outperforms state-of-the-art methods in long-range graph prediction tasks, with a fraction of the computational cost in both FLOPs and GPU memory consumption. The code and models are publicly available at https://github.com/bowang-lab/Graph-Mamba.

연구 동기 및 목표

• 대형 그래프에서의 장거리 의존성에 대한 Graph Transformer의 계산 병목 현상을 동기부여하고 해결한다.
• 입력 의존적 노드 선택 및 컨텍스트 필터링을 수행하는 그래프에 맞춘 선택적 상태 공간 모델링(GMB) 모듈을 개발한다.
• GMB를 GraphGPS 프레임워크에 통합하여 Graph-Mamba를 만들고 비순차적 그래프 데이터에서의 효율적인 장거리 추론을 가능하게 한다.
• 그래프 특화 시퀀스 모델링의 안정성과 성능 향상을 위한 노드 우선순위 지정 및 순열 기반 학습/추론 레시피를 제안한다.

제안 방법

• 그래프에 그래프 중심의 노드 우선순위 지정 및 순열 전략을 도입하여 선택적 상태 공간 모델인 Mamba를 그래프에 적응시킨다.
• GraphGPS의 주의 메커니즘을 Graph-Mamba Block(GMB)으로 대체하여 MPNN 기반의 로컬 컨텍스트와 노드 중심의 GMB를 결합한 전역 컨텍스트를 구현한다.
• 입력 의존적 노드 선택과 시퀀스 인식 노드 우선순위를 가능하게 하는 두 단계의 희소화(스파시피케이션) 메커니즘을 구현한다.
• 입력 길이에 선형 시간 복잡도를 달성하기 위해 GPU 친화적인 상태 공간 업데이트의 이산화/매개변수화를 사용한다.
• 노드 순서 간의 순열 불변성과 안정성을 촉진하기 위한 순열 기반 학습 및 추론 레시피를 도입한다.
• Dense Transformer 및 희소 주의 베이스라인과 비교하여 열 개의 Long Range Graph Benchmark 데이터셋에서 평가한다.

실험 결과

연구 질문

• RQ1 그래프에 적응된 선택적 상태 공간 모델(Graph-Mamba)가 계산 비용을 줄이면서도 장거리 의존성 작업에서 밀집/희소 그래프 주의의 예측 능력에 맞추거나 이를 능가할 수 있는가?
• RQ2 입력 의존적 노드 우선순위 지정 및 순열 기반 학습이 비순차적 그래프 데이터에서의 장거리 추론을 개선하는가?
• RQ3 최신 주의 메커니즘과 비교하여 Graph-Mamba가 다양한 그래프 예측 작업(노드, 그래프, 엣지 수준) 및 그래프 크기에 대해 어떻게 성능을 보이는가?

주요 결과

모델	펩타이드-함수 AP↑	펩타이드-구조 MAE↓	PascalVOC-SP F1↑	COCO-SP F1↑	MALNET-TINY 정확도↑
GCN	0.5930 ± 0.0023	0.3496 ± 0.0013	0.1268 ± 0.0060	0.0841 ± 0.0010	0.8100
GIN	0.5498 ± 0.0079	0.3547 ± 0.0045	0.1265 ± 0.0076	0.1339 ± 0.0044	0.8898 ± 0.0055
GatedGCN	0.5864 ± 0.0077	0.3420 ± 0.0013	0.2873 ± 0.0219	0.2641 ± 0.0045	0.9223 ± 0.0065
GPS+Transformer	0.6575 ± 0.0049	0.2510 ± 0.0015	0.3689 ± 0.0131	0.3774 ± 0.0150	OOM (bs=8)
GPS+Performer	0.6475 ± 0.0056	0.2558 ± 0.0012	0.3724 ± 0.0131	0.3761 ± 0.0101	0.9264 ± 0.0078
GPS+BigBird	0.5854 ± 0.0079	0.2842 ± 0.0130	0.2762 ± 0.0069	0.2622 ± 0.0008	0.9234 ± 0.0034
Exphormer	0.6258 ± 0.0092	0.2512 ± 0.0025	0.3446 ± 0.0064	0.3430 ± 0.0108	0.9422 ± 0.0024
Graph-Mamba	0.6739 ± 0.0087	0.2478 ± 0.0016	0.4191 ± 0.0126	0.3960 ± 0.0175	0.9340 ± 0.0027

• Graph-Mamba는 10개 데이터셋에 걸쳐 장거리 그래프 예측 작업에서 최첨단 또는 경쟁력 있는 결과를 달성하며, 특히 150–1,400노드 규모의 그래프에서 두드러진 성능을 보인다.
• 대형 그래프에서 Graph-Mamba는 다른 희소 주의 방법보다 최대 5% 향상을 제공하고 여러 케이스에서 밀집 Transformer 베이스라인을 능가한다.
• Graph-Mamba는 선형 시간 복잡도와 감소된 GPU 메모리 사용을 보여주며, 대형 그래프에서 최대 74%의 메모리 절감, MalNet-Tiny에서 Transformer 대비 66%의 FLOPs 감소를 달성한다.
• 차수에 따른 노드 우선순위 지정과 노드 수준 순열 결합으로 눈에 띄는 성능 향상을 얻고, 기본 학습/추론 레시피로 권장된다.
• 실험적 FLOPs/메모리 벤치마크는 Graph-Mamba가 데이터셋 전반에 걸쳐 밀집 및 희소 주의 베이스라인보다 비용이 낮아 더 큰 그래프에서의 학습을 가능하게 한다.