Graph State Space Models
- Graph State Space Models are statistical frameworks that combine state space models with graph-structured dependencies to capture interactions among latent variables.
- They employ sparse polynomial and linear parameterizations to recover both dynamic interactions and underlying graph topology through regularized likelihood optimization.
- These models achieve superior recovery and interpretability in complex systems by leveraging scalable inference methods such as proximal gradient and EM algorithms.
A Graph State Space Model (GSSM) is a statistical modeling framework that combines classic state space models with graph-structured dependencies among latent system variables. In GSSMs, the evolution of the latent state is explicitly governed by a (typically sparse) interaction graph encoding which components influence each other, supporting both linear and nonlinear dynamics. This formalism enables modelers to impose, infer, or exploit structured dependencies in large-scale or high-dimensional dynamical systems, offering performance and interpretability benefits over models ignoring graph structure (Cox et al., 2024).
1. Mathematical Formulation and Core Principles
A GSSM formalizes the evolution of a latent state vector such that
with observations
Crucially, the transition map and (often) the observation map are constrained by or parameterized through an adjacency matrix , where denotes that component directly influences component .
In GraphGrad, for nonlinear systems, is parameterized as a sparse polynomial: where indexes multi-degrees and ; the sparsity pattern of directly encodes the graph . Influence from to is present when some with and exists (Cox et al., 2024). For linear-Gaussian GSSMs, the transition matrix itself serves as the (weighted, directed) adjacency, and sparsity reflects Granger-causal structure among state components (Chouzenoux et al., 2023, Elvira et al., 2022).
2. Sparse Polynomial and Linear Parameterizations
Approaches to parameterizing and in GSSMs:
- Polynomial Expansion (GraphGrad): The transition is approximated as a sum over monomials up to total degree , with coefficient vectors encoding edge structure. This allows recovery of both graph topology and nonlinear interaction strengths from observed data via -type sparse regularization (Cox et al., 2024).
- Linear State Transition: In linear-Gaussian settings, the state dynamics are , , with sparse and interpreted as an adjacency encoding direct Granger-causal links. Various works, including GraphIT and GraphEM, leverage sparsity-promoting priors (e.g., , non-convex surrogates) and EM-type algorithms to estimate under model and data constraints (Chouzenoux et al., 2023, Elvira et al., 2022).
These parameterizations support estimation of both the graph structure and system dynamics, addressing identifiability and interpretability.
3. Inference and Optimization Techniques
Estimation of GSSM parameters typically relies on penalized likelihood optimization:
- Marginal Likelihood for Nonlinear GSSMs: With polynomial , the negative log-likelihood is approximated by differentiable particle filtering followed by sparsity-penalized optimization:
Stochastic proximal-gradient methods with soft-thresholding yield both sparse and stable solutions (Cox et al., 2024).
- Majorization-Minimization and Proximal Splitting for Linear GSSMs: Both GraphIT and GraphEM develop iterative convex surrogates for the penalized negative log-likelihood, solving at each iteration a convex optimization problem via proximal splitting (e.g., Douglas–Rachford) for or more sophisticated penalties, using EM to alternate between inference and parameter update (Chouzenoux et al., 2023, Elvira et al., 2022).
These frameworks allow for efficient and stable parameter recovery in both dense and highly sparse regimes.
4. Algorithmic Workflows and Computational Aspects
The general optimization workflow in GSSMs includes:
| Stage | Approach (Nonlinear) | Approach (Linear) |
|---|---|---|
| State/Obs. Simulation | Monte Carlo Particle Filter | Kalman Smoother |
| Objective Eval. | Monte Carlo Likelihood | EM Surrogate |
| Regularization | (proximal) penalty | /group/Laplace penalty |
| Optimization | Stochastic Proximal Gradient | Proximal Splitting (DR/etc.) |
For nonlinear polynomial GSSMs, batching and mini-batch optimization (Batched-GraphGrad) are deployed to address likelihood degeneracy in long sequences, while maintaining overall computational efficiency—forward/backward costs are (batches, steps, particles, time) (Cox et al., 2024).
For linear-Gaussian cases, each EM outer-loop involves a Kalman smoother (to generate sufficient statistics) and an inner convex optimization (with per-iteration cost dominated by matrix operations or soft-thresholding), allowing applicability to high-dimensional systems (Chouzenoux et al., 2023, Elvira et al., 2022).
5. Empirical Performance and Recovery Guarantees
Empirical validations span both chaotic nonlinear systems and high-dimensional stochastic linear systems:
- GraphGrad recovers exact support (precision=recall=1.0) and achieves near-oracle RMSE on Lorenz-63 and Lorenz-96 (e.g., for , , perfect for , RMSE ), outperforming dense MLE and naive subgradient approaches (Cox et al., 2024).
- GraphIT/GraphEM accurately reconstruct the true transition matrices and edge supports in synthetic LG-SSMs, with F1-scores and detection accuracy markedly superior to unconstrained EM or naive Granger-causality methods, especially in highly sparse or structured regimes (Chouzenoux et al., 2023, Elvira et al., 2022).
Penalty selection and initialization are critical: proper use of /nonconvex surrogates, majorization-minimization, and variational inference are all empirically necessary for accurate and robust recovery.
6. Extensions, Limitations, and Open Directions
Current GSSM approaches address both linear and nonlinear, directed and undirected (or functional) graphs, but each framework exhibits limits:
- Assumptions: Gaussian noise, separable transition structure, and stationarity are common constraints.
- Extensions: Future directions include error-bound analysis (comparing EKF, Unscented KF, GSSM), automatic identification of constant vs. dynamic variables, non-Gaussian robustness, and scalable real-time inference via incremental solvers or factor graphs (Lü, 2021).
- Limitations: Loopy, multi-connected factor-graph formulations require careful inference (junction tree, window size), and performance may degrade with ill-posed parameterizations or insufficient data.
Ongoing research in GSSMs targets more general dynamics, dynamic graphs, heterogeneous data, and principled methods for structure learning, exploiting advances in sparse optimization and scalable inference (Cox et al., 2024, Chouzenoux et al., 2023, Elvira et al., 2022).
7. Relationship to Broader SSM and Graph Learning Literature
GSSMs generalize classic state space models—Kalman filters, hidden Markov models, polynomial autoregressive models—by incorporating explicit, learnable, or prior-specified graph structure into the dynamics. This positions them at the intersection of dynamical systems, probabilistic graphical models, and recent graph neural network research. Techniques such as particle filtering, proximal-gradient, EM, and variational inference form the algorithmic backbone of modern GSSM estimation (Cox et al., 2024, Chouzenoux et al., 2023, Elvira et al., 2022). The learned or imposed graph structure provides interpretable insights into the system's dynamics and causal structure—an advantage central to applications in neuroscience, multivariate forecasting, physics, and network science.