Message-Passing State-Space Models: Improving Graph Learning with Modern Sequence Modeling

TL;DR

This paper introduces MP-SSM, a novel graph learning framework that integrates modern state-space model principles into message-passing neural networks, enhancing efficiency, interpretability, and long-range information propagation.

Contribution

It presents MP-SSM, a unified, permutation-equivariant graph model that combines SSM principles with message passing, enabling better theoretical analysis and practical performance.

Findings

Efficient and scalable implementation of MP-SSM.

Improved long-range information propagation in graphs.

Strong empirical results across diverse graph tasks.

Abstract

The recent success of State-Space Models (SSMs) in sequence modeling has motivated their adaptation to graph learning, giving rise to Graph State-Space Models (GSSMs). However, existing GSSMs operate by applying SSM modules to sequences extracted from graphs, often compromising core properties such as permutation equivariance, message-passing compatibility, and computational efficiency. In this paper, we introduce a new perspective by embedding the key principles of modern SSM computation directly into the Message-Passing Neural Network framework, resulting in a unified methodology for both static and temporal graphs. Our approach, MP-SSM, enables efficient, permutation-equivariant, and long-range information propagation while preserving the architectural simplicity of message passing. Crucially, MP-SSM enables an exact sensitivity analysis, which we use to theoretically characterize information flow and evaluate issues like vanishing gradients and over-squashing in the deep regime. Furthermore, our design choices allow for a highly optimized parallel implementation akin to modern SSMs. We validate MP-SSM across a wide range of tasks, including node classification, graph property prediction, long-range benchmarks, and spatiotemporal forecasting, demonstrating both its versatility and strong empirical performance.