Message-Passing State-Space Models: Improving Graph Learning with Modern Sequence Modeling
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Graph Regression | Peptides struct LRGB (test) | MAE0.2458 | 238 | |
| Graph Classification | Peptides-func LRGB (test) | AP0.6993 | 196 | |
| Node Classification | Amazon-Ratings (test) | Accuracy53.65 | 155 | |
| Node Classification | Minesweeper (test) | AUROC95.33 | 134 | |
| Node Classification | Tolokers (test) | AUROC85.26 | 128 | |
| Node Classification | Questions (test) | AUROC78.18 | 83 | |
| Spatio-temporal forecasting | Chickenpox Hungary | MSE0.748 | 44 | |
| Spatio-temporal forecasting | PedalMe London | MSE0.647 | 44 | |
| Spatio-temporal forecasting | Wikipedia Math | MSE0.509 | 44 | |
| Traffic Forecasting | METR-LA Horizon 12 | MAE3.17 | 27 |