ST-SSM: Spatiotemporal State Space Neural Operator
- ST-SSM is a framework that combines state-space models with neural operator techniques to map spatiotemporal functions efficiently.
- It leverages Galerkin projections, adaptive SSM blocks, and parallel scans to ensure stability, universality, and linear scaling in complex systems.
- Empirical benchmarks show that ST-SSM achieves lower error rates and faster computation compared to CNNs, Transformers, and other neural operators on various PDE challenges.
A Spatiotemporal State Space Neural Operator (ST-SSM) is a framework designed for learning operators that map spatiotemporal function-valued inputs to outputs, with primary application to modeling and surrogate learning of time-dependent partial differential equations (PDEs), partial integro-differential equations (PIDEs), and high-dimensional dynamical fields under control or external forcing. The core principle is the unification of state-space models (SSMs) for temporal evolution with operator-learning techniques for spatial structure, yielding architectures that integrate memory, stability, high expressivity, and parameter efficiency for complex spatiotemporal systems.
1. Mathematical Foundation and Operator Formulation
ST-SSMs are grounded in an overview of classical control-theoretic state-space models and neural operator paradigms. The general form encompasses both continuous-time and discrete-time updates:
Continuous-time SSM: where is the hidden state, the system input (often a spatial field), the output, and are learned or input-dependent system matrices/functions (Koren et al., 31 Jul 2025, Cheng et al., 2024).
In spatially extended systems, the input is a field with denoting spatial coordinates and time. A Galerkin-based projection is often used to map to a finite-dimensional latent state: 0 where 1 are neural-network-parameterized spatial basis functions (Ruiter et al., 18 Mar 2026).
In state-space neural operator frameworks, the operator learning goal is posed as direct approximation of the solution (or flow map) for PDEs/PIDEs: 2 for systems such as
3
with 4 a nonlinear operator (e.g., nonlocal kernel or neural field term) and 5 the control/exogenous input (Ruiter et al., 18 Mar 2026).
2. Architectural Components and Model Instantiations
ST-SSM architectures can be largely decomposed into the following classes:
- Galerkin-projected neural ODE/PIDE solvers: These learn both the spatial basis (via a neural network) and the reduced-order temporal ODE system, often bypassing explicit integration by directly parameterizing the flow map as a function of initial state and projected inputs (Ruiter et al., 18 Mar 2026).
- Structured spatiotemporal scans: In models such as Mamba Neural Operator and canonical ST-SSMs, structured SSMs are applied along temporal and spatial axes, with factorized spatial SSMs handling distinct dimensions (via bidirectional sweeps or spatial SSM blocks) and temporal SSMs implementing memory and non-Markovian dynamics (Koren et al., 31 Jul 2025, Cheng et al., 2024).
- Selective/Adaptive SSM blocks: ST-Mamba, STG-Mamba, and FR-Mamba instantiate SSM cores whose transition, input, and output kernels are input-dependent and may be gated or fused with graph neural network (GNN) or Fourier-operator branches, often exploiting per-step or per-feature dynamic adaptation for efficiency and stability (Shao et al., 2024, Li et al., 2024, Long et al., 21 May 2025).
- Biological or neuro-dynamical extensions: The Parallelized Hierarchical Connectome (PHC) framework extends SSMs to spiking, hierarchical, and biophysically-constrained domains by separating neuron-intrinsic diagonal recurrences from spatial synaptic interactions and supporting advanced neural constraints such as Dale’s Law, short-term plasticity, and reward-modulated STDP (Chiang, 1 Apr 2026).
3. Training Methodologies and Flow Learning
Many ST-SSM frameworks employ two-phase or composite losses:
- Basis function learning (POD-supervision): For architectures using modal projections, spatial bases are learned to approximate empirical principal components or dominant singular vectors of snapshot matrices, using orthonormality and basis-approximation losses (Ruiter et al., 18 Mar 2026).
- Operator (flow map) learning: The projected state evolution (often an ODE or RNN) is not integrated numerically, but mapped directly via neural nets (e.g., RNNs, LSTMs, Mamba modules) that learn the flow map from initial coefficients and input projections, with loss computed as mean squared error between predicted and true spatiotemporal fields (Ruiter et al., 18 Mar 2026, Koren et al., 31 Jul 2025, Cheng et al., 2024).
- Input-dependent parameterization: Input-adaptive gates and SSM kernels, as in ST-Mamba and STG-Mamba, allow dynamic adaptation of step sizes, transition matrices, and output kernels, often employing softplus nonlinearities or learned Kalman-fusion for multi-scale data (Shao et al., 2024, Li et al., 2024).
- Parallel scan and complexity advantages: Modern diagonal/parallel SSM blocks exploit scan algorithms (with O(log T) parallel depth) for temporal and/or spatial recurrence, yielding efficient implementations with linear complexity in sequence or field size, avoiding quadratic complexity of attention and other token-mixing mechanisms (Smith et al., 2023, Chiang, 1 Apr 2026).
4. Theoretical Properties: Universality, Stability, and Expressivity
A distinguishing feature of the ST-SSM paradigm is rigorous theoretical support:
- Universality: Convolutional neural operators with SSM-based kernels are universal approximators of continuous spatiotemporal operators, provided their full field-of-view criterion is satisfied (i.e., spatial kernels can propagate information globally) (Koren et al., 31 Jul 2025). The Stone–Weierstrass theorem guarantees that such architectures can approximate any continuous operator on compact sets.
- Stability and Causality: By leveraging the state-space structure, ST-SSMs inherit control-theoretic properties. Eigenvalue spectra of A matrices may be regularized (spectral radius or contractivity) to ensure stability, and causality is structurally enforced via recurrence and the form of temporal convolution kernels (e.g., lower-triangular or convolutional in time) (Cheng et al., 2024).
- Interpretation as Neural Operators: Learned Green’s functions in SSMs provide a principled and expressive mechanism for instantiating convolutional integral operators in both space and time (continuous or discrete), generalizing over classical Fourier neural operators, and enabling adaptive locality or globality (Koren et al., 31 Jul 2025, Cheng et al., 2024, Ruiter et al., 18 Mar 2026).
5. Empirical Benchmarks, Comparison, and Efficiency
ST-SSM models demonstrate strong empirical performance and computational efficiency across domains. Key results include:
- PDE Benchmarks: On tasks such as 1D Kuramoto–Sivashinsky, viscous Burgers, and 2D Navier–Stokes, ST-SSM outperforms FNOs, S4-FFNO, and CNN baselines in normalized 6 relative error while using 7–8 fewer parameters. For instance, at 9 spatial resolution, ST-SSM achieves 0 error (vs 1 S4-FFNO) on KS (2) and 3 vs 4 (FNO2D) on 2D Navier–Stokes (Koren et al., 31 Jul 2025).
- Efficiency and Scaling: ST-SSM enables linear scaling in sequence or field size, substantial speedups over Transformer-based attention (e.g., 61% faster than SOTA Transformer on traffic flow, matching or exceeding accuracy), and constant-time autoregressive prediction desirable for physical forecasting and real-time control (Shao et al., 2024, Smith et al., 2023).
- Transfer Learning and Adaptability: In control-theoretic operator frameworks such as RHYME-XT, partial fine-tuning on novel data (e.g., new kernels in neural field equations) yields fast convergence and accuracy with minimal retraining, supporting generalization across data regimes (Ruiter et al., 18 Mar 2026).
6. Extensions, Ablations, and Practical Considerations
- Spatial Factorization and Bidirectionality: Empirical ablations highlight the necessity of bidirectional spatial sweeps for universal expressivity and sharp error optima. Unidirectional (causal) spatial SSMs are non-universal and plateau in error (Koren et al., 31 Jul 2025).
- Role of Temporal Memory: Under conditions of partial observability or unknown system parameters, temporal memory in the SSM acts as a sufficient statistic, inferring missing context and maintaining performance robustness (Koren et al., 31 Jul 2025).
- Adaptive and Hybrid Blocks: Several instantiations combine SSMs with Fourier (FNO), convolutional operators, or graph neural modules (e.g., STG-Mamba’s fusion with Kalman-filtered GNN features) to cover heterogeneous or multi-modal spatiotemporal data (Li et al., 2024, Long et al., 21 May 2025).
- Biological Constraints: PHC-style ST-SSMs show that complex bio-physical priors—including leaky integrate-and-fire, short-term plasticity, and reward modulation—can be imposed without loss of efficiency, suggesting a principled path to parameter-efficient and biologically plausible sequence models (Chiang, 1 Apr 2026).
7. Summary Table: Model Attributes and Experimental Regimes
| Architecture | Spatiotemporal Mix | Complexity | Empirical Regimes |
|---|---|---|---|
| RHYME-XT | Learned Galerkin flow | Two-stage, projection + flow | Neural field PIDEs |
| MNO/ST-SSM [Mamba] | Diagonal/parallel SSM scan | O(L), linear in field size | Standard PDE benchmarks, data-scarcity |
| ST-Mamba | Selective state/adaptivity | O(TN), linear in nodes | Traffic flow forecasting |
| STG-Mamba | Kalman-GNN adaptivity + SSM | O(T), linear | Spatial-temporal graph, system embedding |
| ConvSSM/ConvS5 | Convolutional SSM | O(L), subquadratic | Long-range video, AR environments |
| PHC/PHCSSM | Diagonal SSM + multi-transmission | O(log T), biologically plausible | Physiological time series |
Each instantiation validates the general principle that spatiotemporal state-space operators offer an efficient, robust, and theoretically principled architecture for learning complex solution maps of modern time-dependent systems (Ruiter et al., 18 Mar 2026, Cheng et al., 2024, Koren et al., 31 Jul 2025, Smith et al., 2023, Shao et al., 2024, Li et al., 2024, Chiang, 1 Apr 2026, Long et al., 21 May 2025).