Technical Summary: Message-Passing State-Space Models (MP-SSM)

Problem Statement

Graph Neural Networks (GNNs), particularly Message-Passing Neural Networks (MPNNs), are standard for learning from graph-structured data. However, traditional MPNNs (e.g., GCNs) struggle with long-range dependencies due to over-squashing (the compression of information from distant nodes) and vanishing gradients. While various remedies exist—such as graph rewiring, Transformers, or spectral methods—many suffer from high computational costs, lack of permutation equivariance, or reliance on sequential graph representations that distort topology.

Concurrently, State-Space Models (SSMs) like Mamba and S4 have achieved success in sequence modeling by utilizing linear recurrent dynamics to ensure stable gradient flow and efficient long-range dependency modeling. Recent attempts to adapt SSMs to graphs (Graph SSMs) often fail to integrate seamlessly with the message-passing paradigm, relying on spectral decompositions or sequentializations that compromise theoretical guarantees on information propagation and stability.

Crucially, existing Graph SSM approaches lack rigorous, exact characterizations of how information flows between nodes, offering only loose or qualitative guarantees regarding sensitivity and stability in deep architectures.

Methodology

The authors propose the Message-Passing State-Space Model (MP-SSM), a framework that integrates modern SSM computation directly into the MPNN paradigm through the lens of sensitivity analysis.

Core Architecture

MP-SSM unifies the treatment of static and temporal graphs using a linear state-space recurrence defined on the graph shift operator (GSO).

• Graph Representation: The graph $G=(V, E)$ is processed using a symmetrically normalized adjacency matrix with self-loops ($A$) as the GSO.
• Linear Recurrence: Unlike standard MPNNs that apply nonlinearities at every step, MP-SSM employs a purely linear recurrence for the state update:
  $$X_{t+1} = A X_t W + U_{t+1} B$$
  where $X_t$ represents node embeddings, $U_t$ are input features, and $W, B$ are learnable weight matrices.
• Nonlinearity Placement: Nonlinearities are restricted to a Multi-Layer Perceptron (MLP) decoder applied after the recurrence steps (or at the final step for static graphs). This design choice is critical for enabling exact sensitivity analysis and efficient parallel implementation.
• Static vs. Temporal:
  • Temporal: The model processes a sequence of graphs $[U_1, \dots, U_{k+1}]$ naturally.
  • Static: To unify the framework, a single input $U_1$ is treated as a constant sequence $[U_1, \dots, U_1]$. The recurrence effectively aggregates information over $k$ hops, and the MLP is applied to the final state.
• Deep Architecture: Stacking multiple MP-SSM blocks allows the model to aggregate information from $s \times k$ hops, where $s$ is the number of blocks and $k$ is the recurrence depth.

Implementation

The authors derive a closed-form solution for the recurrence by unrolling the linear equation. This allows for a fast parallel implementation (Algorithm 1) that computes the entire sequence in a single step, avoiding the sequential bottlenecks of standard RNNs. The implementation leverages the diagonalization of the GSO and weight matrices to compute powers efficiently, achieving speedups comparable to standard GCNs while supporting much deeper propagation.

Key Contributions

1. Principled Integration of SSMs and MPNNs:
   The paper introduces MP-SSM, a framework that embeds linear state-space dynamics directly into message passing. This design preserves permutation equivariance, maintains the original graph topology, and supports both static and temporal graphs without sequentialization or spectral approximations.

2. Exact Sensitivity Analysis:
   A primary theoretical contribution is the derivation of closed-form Jacobian computations for node-to-node dependencies.

   • Local Sensitivity: The authors prove that the Jacobian between node $j$ at step $s$ and node $i$ at step $t$ is exactly $(A^{t-s})_{ij} (W^\top)^{t-s}$.
   • Stability Guarantees: Using the properties of the symmetrically normalized adjacency matrix, they establish that the spectral radius is bounded by 1, ensuring stable dynamics.
   • Lower Bounds: The paper derives precise lower bounds for both local sensitivity (addressing over-squashing) and global sensitivity (addressing vanishing gradients). Specifically, Theorem 3.11 proves that global sensitivity is strictly greater than zero for connected graphs, a guarantee not provided by standard nonlinear MPNNs where nonlinearity derivatives can cause information to vanish.
   • Comparison to GCN: Theorem 3.13 demonstrates that a standard GCN with ReLU activations vanishes information at a rate of $2^{-k/2}$ faster than MP-SSM for $k$ steps.

3. Empirical Validation:
   MP-SSM is evaluated across 15 benchmarks, including synthetic long-range tasks, heterophilic node classification, and spatiotemporal forecasting. The model consistently matches or outperforms state-of-the-art baselines (including Graph Transformers, rewired GNNs, and other Graph SSMs) while maintaining computational efficiency comparable to standard GCNs.

Experimental Results

• Graph Property Prediction: On synthetic tasks requiring long-range flow (Diameter, SSSP, Eccentricity), MP-SSM significantly outperforms all baselines, including A-DGN and GCNII, with average gains of 2.4 points in log(MSE).
• Heterophilic Benchmarks: On five heterophilic datasets (e.g., Roman-empire, Amazon-ratings), MP-SSM achieves the highest average ranking, improving over GCN by up to 17% and surpassing specialized heterophily models and Transformers.
• Spatiotemporal Forecasting: On traffic (Metr-LA, PeMS-Bay) and other time-series datasets, MP-SSM outperforms existing temporal GNNs and dedicated state-space approaches like GGRNN and GraphSSM-S4.
• Long-Range Graph Benchmark (LRGB): On Peptides-func and Peptides-struct, MP-SSM achieves top-tier performance without relying on global attention mechanisms or graph rewiring, demonstrating its ability to capture structural information through linear recurrence alone.
• Efficiency: Runtime analysis shows MP-SSM scales linearly with graph size, comparable to GCNs, and significantly faster than Graph Transformers or deep MPNNs with many layers. The fast parallel implementation further reduces inference time.

Significance and Claims

The paper claims that sensitivity analysis provides a principled foundation for designing Graph State-Space Models. By removing nonlinearities from the recurrent dynamics and placing them in a decoder, MP-SSM achieves:

1. Theoretical Rigor: Exact guarantees on information flow, proving that the model avoids vanishing gradients and over-squashing under specific conditions, unlike heuristic approaches.
2. Architectural Simplicity: A unified framework that handles both static and temporal graphs without complex modifications like rewiring or spectral decomposition.
3. Practical Effectiveness: Strong empirical performance across diverse tasks, demonstrating that linear recurrence is sufficient for capturing long-range dependencies when combined with appropriate deep learning heuristics (residual connections, normalization) between blocks.

The authors conclude that MP-SSM represents a shift from heuristic adaptations of SSMs to graphs toward a mathematically grounded integration that preserves the core benefits of both message passing and state-space modeling.