Multivariate low-rank state–space model with SPDE approach for high-dimensional data

This paper proposes a novel low-rank approximation of the State–Space Model (SSM) with spatially correlated innovations for the analysis of multivariate spatio-temporal data. The SSM’s measurement equation is based on a linear coregionalisation model, which describes the cross-correlation between the observed variables, while the Matérn Gaussian innovation term in the state equation is modelled using the Stochastic Partial Differential Equation (SPDE) approach, allowing a finite-dimensional representation of the latent processes using basis functions defined on spatial meshes. Dimensionality reduction is achieved by appropriately reducing the number of nodes in the meshes. Inference on the model parameters is performed via Maximum Likelihood Estimation (MLE), implemented through the Expectation–Maximisation (EM) algorithm, which features closed-form updating formulas for most parameters and efficient numerical routines for the remainder. We derive theoretical results on the accuracy and convergence of the low-rank approximation and validate them through simulation studies. The EM algorithm and the likelihood derivatives required for inference are implemented in Python/JAX, enabling automatic differentiation and scalable execution across all available local CPU cores, with native support for GPU and TPU acceleration. By analysing a large bivariate air-quality dataset, we demonstrate that reducing the number of nodes by 75% enables model estimation to be 15.8 times faster with only a 15% increase in validation error. We also compare our approach with SPDE-INLA alternatives, demonstrating improved computational scalability while maintaining comparable predictive performance.