Computation-Aware Kalman Filtering and Smoothing

TL;DR

The paper tackles the prohibitive cost of inference in high-dimensional linear-Gaussian state-space models by introducing computation-aware Kalman filtering and smoothing (CAKF/CAKS). These methods use matrix-free, iterative updates with low-dimensional observation projections and downdate truncation to dramatically reduce time and memory while embedding the approximation error into the posterior uncertainty. Theoretical results provide complexity bounds and a pointwise error bound showing the CAKS uncertainty upper-bounds prediction error, and experiments on synthetic data and a large ERA5 climate dataset demonstrate superior performance and scalability relative to ensemble methods and standard KF/RTS. This work offers a practical route to scalable, uncertainty-aware spatiotemporal regression and GP-style inference in very large state spaces, leveraging GPU acceleration and probabilistic numerics principles.

Abstract

Kalman filtering and smoothing are the foundational mechanisms for efficient inference in Gauss-Markov models. However, their time and memory complexities scale prohibitively with the size of the state space. This is particularly problematic in spatiotemporal regression problems, where the state dimension scales with the number of spatial observations. Existing approximate frameworks leverage low-rank approximations of the covariance matrix. But since they do not model the error introduced by the computational approximation, their predictive uncertainty estimates can be overly optimistic. In this work, we propose a probabilistic numerical method for inference in high-dimensional Gauss-Markov models which mitigates these scaling issues. Our matrix-free iterative algorithm leverages GPU acceleration and crucially enables a tunable trade-off between computational cost and predictive uncertainty. Finally, we demonstrate the scalability of our method on a large-scale climate dataset.

Figures (11)

• Figure 1: $D = 14640.0$
• Figure 2: $D = 231360.0$
• Figure 4: Comparison of the CAKF, the EnKF, and two variants of the ETKF on on-model data with state-space dimension $D = 20000.0$ while varying the rank parameters that govern the computational budget of the algorithms. The CAKF significantly outperforms the other filter variants for high-dimensional state spaces.
• Figure 5: Predictive mean, predictive standard deviation, and pointwise absolute error for an increasing maximal number of iterations $\check{N}^{\textnormal{max}} \ge \check{N}_k$ per time step on a synthetic spatiotemporal regression problem.
• Figure 6: Work-precision diagrams for the CAKF and CAKS on the ERA5 climate dataset. The plot shows the mean squared error (MSE) and average negative log density (NLD) of the computation-aware filter and smoother for different problem sizes (i.e., state-space dimension) and number of iterations on the train and test set. The predictive error measured by test MSE decreases with larger problem sizes, while the test NLD increases. This is because we limit the computational budget and thus run fewer iterations for larger problems, i.e., we trade reduced computation cost for increased uncertainty.

Theorems & Definitions (34)

• theorem 1: Pointwise Worst-Case Prediction Error
• definition B.1: Linear-Gaussian State-Space Model
• theorem B.2: Kalman Filter
• proposition B.3: Downdate-Form Kalman Filter
• proof
• theorem B.4: RTS Smoother
• proposition B.5: Inverse-Free RTS Smoother
• proof
• lemma B.6: Matheron's Rule
• proof