Computation-Aware State-Space Model (CASSM)

• CASSM is a probabilistic framework for high-dimensional state-space models that combines low-rank projections and iterative solvers to enable tractable Bayesian inference.
• It reduces computational complexity by projecting high-dimensional observations into a lower-dimensional space and updating covariances via matrix-free methods.
• CASSM is applied in domains like climate modeling and neural data analysis, providing calibrated uncertainty estimates and efficient model learning.

The Computation-Aware State-Space Model (CASSM) provides a scalable, uncertainty-calibrated probabilistic framework for filtering, smoothing, and learning in high-dimensional state-space models (SSMs). Designed for applications where state or observation dimensions reach into the tens or hundreds of thousands, CASSM combines low-rank matrix approximations, iterative solvers, and explicit modeling of computational uncertainty to enable tractable and rigorous Bayesian inference. CASSM can be applied to spatiotemporal processes in climate modeling, high-dimensional neural data, and other regimes where classical Kalman filters are intractable due to cubic complexity and memory bottlenecks (Pförtner et al., 2024, Huml et al., 31 May 2026).

1. Foundational Gauss–Markov State–Space Model

CASSM is built upon the standard linear–Gaussian state-space model (LGSSM), where the latent state $x_k \in \mathbb{R}^d$ evolves as

$$x_0 \sim \mathcal{N}(m_0, P_0), \quad x_k = A_k x_{k-1} + w_{k-1}, \quad w_{k-1} \sim \mathcal{N}(0, Q_{k-1}),$$

and observations $y_k \in \mathbb{R}^{n_k}$ (or $y_k \in \mathbb{R}^N$ for neural data) follow

$$y_k = H_k x_k + v_k, \quad v_k \sim \mathcal{N}(0, R_k).$$

For neural dynamics, the model uses $A, Q, C, R$ (with $C$ as the loading matrix), and is typically initialized with Gaussian priors for $x_0$.

The Kalman filter and Rauch–Tung–Striebel (RTS) smoother provide exact filtering and smoothing recursions, but their $O(n_T d^3)$ time and $O(d^2)$ memory complexity are intractable for large $$x_0 \sim \mathcal{N}(m_0, P_0), \quad x_k = A_k x_{k-1} + w_{k-1}, \quad w_{k-1} \sim \mathcal{N}(0, Q_{k-1}),$$0 or $$x_0 \sim \mathcal{N}(m_0, P_0), \quad x_k = A_k x_{k-1} + w_{k-1}, \quad w_{k-1} \sim \mathcal{N}(0, Q_{k-1}),$$1. As a result, approximate methods often sacrifice fidelity or underestimate uncertainty due to computational approximations (Pförtner et al., 2024, Huml et al., 31 May 2026).

2. Computation-Aware Approximations

CASSM mitigates scalability bottlenecks by combining two principal modifications:

• Low-Dimensional Actions/Projections: At each step, a matrix $$x_0 \sim \mathcal{N}(m_0, P_0), \quad x_k = A_k x_{k-1} + w_{k-1}, \quad w_{k-1} \sim \mathcal{N}(0, Q_{k-1}),$$2 or $$x_0 \sim \mathcal{N}(m_0, P_0), \quad x_k = A_k x_{k-1} + w_{k-1}, \quad w_{k-1} \sim \mathcal{N}(0, Q_{k-1}),$$3 (for neural data, $$x_0 \sim \mathcal{N}(m_0, P_0), \quad x_k = A_k x_{k-1} + w_{k-1}, \quad w_{k-1} \sim \mathcal{N}(0, Q_{k-1}),$$4) projects observations into a lower-dimensional space. The projected observation model yields

$$x_0 \sim \mathcal{N}(m_0, P_0), \quad x_k = A_k x_{k-1} + w_{k-1}, \quad w_{k-1} \sim \mathcal{N}(0, Q_{k-1}),$$5

with $$x_0 \sim \mathcal{N}(m_0, P_0), \quad x_k = A_k x_{k-1} + w_{k-1}, \quad w_{k-1} \sim \mathcal{N}(0, Q_{k-1}),$$6 and $$x_0 \sim \mathcal{N}(m_0, P_0), \quad x_k = A_k x_{k-1} + w_{k-1}, \quad w_{k-1} \sim \mathcal{N}(0, Q_{k-1}),$$7. The Kalman update is performed in this subspace, reducing the inversion cost for the innovation covariance from $$x_0 \sim \mathcal{N}(m_0, P_0), \quad x_k = A_k x_{k-1} + w_{k-1}, \quad w_{k-1} \sim \mathcal{N}(0, Q_{k-1}),$$8 to $$x_0 \sim \mathcal{N}(m_0, P_0), \quad x_k = A_k x_{k-1} + w_{k-1}, \quad w_{k-1} \sim \mathcal{N}(0, Q_{k-1}),$$9.

• Iterative, Matrix-Free Covariance Representation: The state covariance is updated as a sequence of low-rank downdates without explicitly forming or storing the full $y_k \in \mathbb{R}^{n_k}$0 matrix:

$y_k \in \mathbb{R}^{n_k}$1

where $y_k \in \mathbb{R}^{n_k}$2 collects downdate directions (computed using conjugate gradients or Lanczos iterations). The application of $y_k \in \mathbb{R}^{n_k}$3 involves only matvecs, making the procedure amenable to GPU acceleration (Pförtner et al., 2024).

If the number of downdate directions grows large, a truncated SVD keeps only the top $y_k \in \mathbb{R}^{n_k}$4 directions, with the remainder incorporated into an additive computational noise term:

$y_k \in \mathbb{R}^{n_k}$5

thus ensuring conservative uncertainty inflation (Pförtner et al., 2024).

3. Modeling and Quantifying Computational Uncertainty

CASSM explicitly quantifies the uncertainty induced by computational approximations. The policy matrices $y_k \in \mathbb{R}^{n_k}$6 are treated as hyperparameters/contextual actions, and the approximate posterior $y_k \in \mathbb{R}^{n_k}$7 is penalized for divergence from the true predictive prior $y_k \in \mathbb{R}^{n_k}$8.

Training minimizes a computation-aware evidence lower bound (ELBO):

$y_k \in \mathbb{R}^{n_k}$9

where $y_k \in \mathbb{R}^N$0 uses the full $y_k \in \mathbb{R}^N$1-dimensional (unprojected) likelihood. A compact numerical form involves only the approximate covariances of size $y_k \in \mathbb{R}^N$2 and $y_k \in \mathbb{R}^N$3 (Huml et al., 31 May 2026).

This ELBO structure ensures that both data fit and the divergence from the exact sequential prior are controlled, and the posterior variance reflects data noise and computational error, thereby avoiding overconfident inference commonly termed "variance starvation."

4. Inference Algorithms, Learning, and Complexity

The main filtering and smoothing operations in CASSM utilize matrix-free, GPU-accelerated iterative solvers. Key steps include:

• Filter Step: At each $y_k \in \mathbb{R}^N$4, the projected data and low-dimensional innovation are computed, followed by iterative conjugate-gradients or Lanczos solves to construct downdate directions for the covariance. Orthogonalization and SVD truncation are used as necessary.
• GPU Acceleration: Custom CUDA kernels enable efficient Gram matrix-vector products. Batched operations and low-level kernel routines are used for orthogonalization and Lanczos tridiagonalization (Pförtner et al., 2024).
• Parameter Learning: All model parameters and projection policies $y_k \in \mathbb{R}^N$5 are trained end-to-end via gradient methods (Adam). Differentiable SVDs circumvent instability in gradient computation through truncation steps (Huml et al., 31 May 2026).

The computational complexity per step is summarized as follows:

Operation	Standard Kalman	CASSM
Covariance storage	$y_k \in \mathbb{R}^N$6	$y_k \in \mathbb{R}^N$7 (downdates, $y_k \in \mathbb{R}^N$8)
Innovation inversion	$y_k \in \mathbb{R}^N$9	$$y_k = H_k x_k + v_k, \quad v_k \sim \mathcal{N}(0, R_k).$$0 ($$y_k = H_k x_k + v_k, \quad v_k \sim \mathcal{N}(0, R_k).$$1)
Overall time (per step)	$$y_k = H_k x_k + v_k, \quad v_k \sim \mathcal{N}(0, R_k).$$2	$$y_k = H_k x_k + v_k, \quad v_k \sim \mathcal{N}(0, R_k).$$3 or $$y_k = H_k x_k + v_k, \quad v_k \sim \mathcal{N}(0, R_k).$$4
Overall memory	$$y_k = H_k x_k + v_k, \quad v_k \sim \mathcal{N}(0, R_k).$$5	$$y_k = H_k x_k + v_k, \quad v_k \sim \mathcal{N}(0, R_k).$$6

If $$y_k = H_k x_k + v_k, \quad v_k \sim \mathcal{N}(0, R_k).$$7 and $$y_k = H_k x_k + v_k, \quad v_k \sim \mathcal{N}(0, R_k).$$8 are small and kernel matvecs are $$y_k = H_k x_k + v_k, \quad v_k \sim \mathcal{N}(0, R_k).$$9 (or $A, Q, C, R$0), the total cost is $A, Q, C, R$1 time and $A, Q, C, R$2 memory, linear in $A, Q, C, R$3 and favorable relative to cubic/quadratic baseline (Pförtner et al., 2024).

5. Model Selection, Learnable Projections, and Policies

In CASSM, projection matrices $A, Q, C, R$4 are not fixed but are treated as tunable policies:

• Entropy-Based Policy: Theoretically, the optimal $A, Q, C, R$5 for minimizing posterior entropy projects onto the top $A, Q, C, R$6 eigenvectors of the innovation covariance $A, Q, C, R$7 [(Huml et al., 31 May 2026), Theorem 1]. However, direct computation is intractable for large $A, Q, C, R$8, so $A, Q, C, R$9 are parameterized (block-sparse, or structured) and learned end-to-end.
• Greedy and End-to-End Training: The parameterization supports backpropagation through time, and block-sparsity keeps the number of tunable parameters $C$0. Differentiable SVD algorithms ensure stable gradient flow through the low-rank truncation process.
• Flexibility: The policy selection can use coordinate actions, randomization, Bayesian-optimal variants, or conjugate-gradient directions. The iterative algorithm allows for per-step or global budget control of the compression parameter $C$1 (Pförtner et al., 2024).

This learnable subspace approach provides effective model-selection in the context of high-dimensional neural or spatiotemporal data when classical model selection is computationally unfeasible.

6. Empirical Scaling, Calibration Guarantees, and Extensions

CASSM’s key theoretical performance and empirical findings are:

• Calibration Guarantee: For RKHS-based spatiotemporal GP regression, the computation-aware smoother posterior mean and variance satisfy the exact worst-case calibration bound as the true posterior, i.e.,

$C$2

thus the posterior uncertainty envelopes both data and computational errors (Pförtner et al., 2024).

• Empirical Results on Large-Scale Problems:
  • On large-scale climate data ($C$3, $C$4), CASSM achieves exponential improvement in MSE and NLL versus the number of CG steps $C$5, outperforming coordinate/random action baselines by orders of magnitude.
  • In neural dynamics benchmarks (e.g., Lorenz attractor, primate datasets, zebrafish recordings), CASSM matches or outperforms deep sequence models (e.g., LFADS) in the low-trials/high-neurons regime, and delivers well-calibrated uncertainty, with 95% coverage between 0.90–0.92 compared to $C$60.7 for baseline GPFA models.
  • On N = 8.8×10⁴ whole-brain zebrafish data, CASSM achieves MSE and NLL matching or surpassing LFADS, with GPFA unable to run due to memory (Huml et al., 31 May 2026).

A notable property is that truncation only inflates the posterior variance—never reducing filter mean accuracy—ensuring conservative inference.

CASSM is implemented in Julia (ComputationAwareKalman.jl) with custom GPU kernels, supports plug-and-play for SSMs with fast matvecs (e.g., PDE solvers, random Fourier features), and posterior sampling via Matheron's rule within $C$7 cost. Likelihood-based learning (e.g., EM) is possible, but gradient computation becomes involved when truncation is used (Pförtner et al., 2024).

7. Applications and Domain Relevance

CASSM is suited for domains demanding tractable, rigorous Bayesian inference in large latent or observed spaces:

• Spatiotemporal Gaussian Process Regression: Efficient for spatial grids with millions of points, such as climate datasets.
• Neural Data Analysis: Particularly effective in the scale-imbalanced regime, where the number of neurons ($C$8) far exceeds the number of trials ($C$9), avoiding overfitting and under-calibration typical in overparameterized deep models (Huml et al., 31 May 2026).
• Control and PDE Surrogates: Models with structured kernel covariance (e.g., Matérn, banded, or PDE-induced covariance) can exploit the CASSM framework for scalable, uncertainty-aware inference.

A plausible implication is that, by explicitly encoding computational error, downstream applications—for example, control design or policy learning—avoid the common pitfall of overconfident, unreliable uncertainties encountered when using purely approximate Bayesian filtering or smoothing without computation-aware mechanisms.

---

References:

(Pförtner et al., 2024) "Computation-Aware Kalman Filtering and Smoothing" (Huml et al., 31 May 2026) "Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics"