What Can We Learn from State Space Models for Machine Learning on Graphs?

TL;DR

This paper introduces Graph State Space Convolution (GSSC), a novel model that extends state space models to graph data, offering improved expressiveness and efficiency over traditional message passing neural networks and transformers.

Contribution

We propose GSSC, a new graph neural network model that leverages state space models with permutation-equivariant set aggregation and relative node distances, enhancing expressiveness and scalability.

Findings

GSSC outperforms baselines on 6 of 11 benchmark datasets.

GSSC demonstrates provably stronger expressiveness than MPNNs.

GSSC achieves significant improvements in capturing long-range dependencies.

Abstract

Machine learning on graphs has recently found extensive applications across domains. However, the commonly used Message Passing Neural Networks (MPNNs) suffer from limited expressive power and struggle to capture long-range dependencies. Graph transformers offer a strong alternative due to their global attention mechanism, but they come with great computational overheads, especially for large graphs. In recent years, State Space Models (SSMs) have emerged as a compelling approach to replace full attention in transformers to model sequential data. It blends the strengths of RNNs and CNNs, offering a) efficient computation, b) the ability to capture long-range dependencies, and c) good generalization across sequences of various lengths. However, extending SSMs to graph-structured data presents unique challenges due to the lack of canonical node ordering in graphs. In this work, we propose Graph State Space Convolution (GSSC) as a principled extension of SSMs to graph-structured data. By leveraging global permutation-equivariant set aggregation and factorizable graph kernels that rely on relative node distances as the convolution kernels, GSSC preserves all three advantages of SSMs. We demonstrate the provably stronger expressiveness of GSSC than MPNNs in counting graph substructures and show its effectiveness across 11 real-world, widely used benchmark datasets. GSSC achieves the best results on 6 out of 11 datasets with all significant improvements compared to the state-of-the-art baselines and second-best results on the other 5 datasets. Our findings highlight the potential of GSSC as a powerful and scalable model for graph machine learning. Our code is available at https://github.com/Graph-COM/GSSC.