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

2025-05-24

AI Summary

Existing graph state space models (GSSMs) map graphs to sequences for sequential SSM processing, thereby violating permutation equivariance, impairing message-passing compatibility, and reducing computational efficiency. This paper proposes MP-SSM—the first model to natively integrate the core state space mechanism into the message-passing framework—yielding a unified, permutation-equivariant graph learning architecture that naturally supports both static and dynamic graphs. Theoretically, we conduct the first sensitivity analysis of gradient vanishing and over-compression phenomena in graph neural networks. Architecturally, we design a parallelizable, efficient tensor-based computation scheme. Experiments demonstrate that MP-SSM consistently outperforms state-of-the-art GNNs and GSSMs across diverse tasks—including node classification, graph property prediction, long-range benchmarks, and spatiotemporal forecasting—achieving superior generalization and empirical performance.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Reasoning under Uncertainty: Graphical Models
• Data Mining & Knowledge Management: Graph Mining, Social Network Analysis & Community

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Enhancing graph learning with modern sequence modeling techniques
• Ensuring permutation equivariance and message-passing compatibility in GSSMs
• Addressing vanishing gradients and over-squashing in deep graph networks

Innovation

• Embedding SSM principles into Message-Passing Neural Networks
• Enabling efficient and permutation-equivariant long-range information propagation
• Allowing exact sensitivity analysis for theoretical characterization