将状态空间模型思想融入消息传递网络,提升图学习的效率与长程建模能力。
Message-Passing State-Space Models: Improving Graph Learning with Modern Sequence Modeling
- 直接在消息传递框架中嵌入现代状态空间模型机制
- 实现高效、可置换等变的长程信息传播,支持深层网络
- 适合需要长程依赖建模的静态与动态图任务
状态空间模型(SSMs)在序列建模中的成功启发了其在图学习中的应用,催生了图状态空间模型(GSSMs)。然而,现有GSSMs通常将SSM模块应用于从图中提取的序列,常牺牲排列等变性、消息传递兼容性和计算效率等核心特性。本文提出一种新视角:将现代SSM计算的核心原理直接嵌入消息传递神经网络框架,构建统一方法用于静态与动态图。所提出的MP-SSM实现高效、排列等变且长程的信息传播,同时保持消息传递的架构简洁性。关键优势在于支持精确敏感性分析,理论上刻画信息流并评估深层网络中的梯度消失与过度压缩问题。此外,设计支持类似现代SSMs的高效并行实现。我们在节点分类、图属性预测、长程基准测试及时空预测等多种任务上验证了MP-SSM,展现出优异的泛化性能与实用性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。