---
title: Computation-Aware State-Space Model (CASSM)
url: https://www.emergentmind.com/topics/computation-aware-state-space-model-cassm
type: topic
---

# Computation-Aware State-Space Model (CASSM)

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 [2405.08971][2606.01468].

## 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 $d$ or $N$. As a result, approximate methods often sacrifice fidelity or underestimate uncertainty due to computational approximations [2405.08971][2606.01468].

## 2. Computation-Aware Approximations

CASSM mitigates scalability bottlenecks by combining two principal modifications:

- **Low-Dimensional Actions/Projections:** At each step, a matrix $U_k \in \mathbb{R}^{n_k \times m_k}$ or $U_k \in \mathbb{R}^{N \times M}$ (for neural data, $M \ll N$) projects observations into a lower-dimensional space. The projected observation model yields
  $$
  \tilde{y}_k = U_k^T y_k \approx (U_k^T H_k) x_k + \tilde{v}_k,
  $$
  with $\tilde{v}_k \sim \mathcal{N}(0, \tilde{R}_k)$ and $\tilde{R}_k = U_k^T R_k U_k$. The Kalman update is performed in this subspace, reducing the inversion cost for the innovation covariance from $O(n_k^3)$ to $O(m_k^3)$.

- **Iterative, Matrix-Free Covariance Representation:** The state covariance is updated as a sequence of low-rank downdates without explicitly forming or storing the full $d \times d$ matrix:
  $$
  P_k = P_k^- - V_k V_k^T,
  $$
  where $V_k \in \mathbb{R}^{d \times m_k}$ collects downdate directions (computed using conjugate gradients or Lanczos iterations). The application of $P_k$ involves only matvecs, making the procedure amenable to GPU acceleration [2405.08971].

If the number of downdate directions grows large, a truncated SVD keeps only the top $r \leq max_k$ directions, with the remainder incorporated into an additive computational noise term:
$$
x_k = A_k x_{k-1} + (w_{k-1} + q_{k-1}^{comp}), \quad q_{k-1}^{comp} \sim \mathcal{N}(0, N_k N_k^T),
$$
thus ensuring conservative uncertainty inflation [2405.08971].

## 3. Modeling and Quantifying Computational Uncertainty

CASSM explicitly quantifies the uncertainty induced by computational approximations. The policy matrices $\{U_k\}$ are treated as hyperparameters/contextual actions, and the approximate posterior $q(x_k \mid y_{1:k}) = \mathcal{N}(\hat{x}_{k|k}, P_{k|k})$ is penalized for divergence from the true predictive prior $p(x_k \mid x_{k-1}) = \mathcal{N}(A \hat{x}_{k-1|k-1}, AP_{k-1|k-1}A^T + Q)$.

Training minimizes a computation-aware evidence lower bound (ELBO):
$$
L(\Theta) = \sum_{k=1}^K \mathbb{E}_{q(x_k|y_{1:k})}[\log p(y_k|x_k)] + D_{KL}(q(x_k|y_{1:k}) \parallel p(x_k|x_{k-1})),
$$
where $p(y_k|x_k)$ uses the full $N$-dimensional (unprojected) likelihood. A compact numerical form involves only the approximate covariances of size $M$ and $d$ [2606.01468].

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 $k$, 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 [2405.08971].
- **Parameter Learning:** All model parameters and projection policies $(A, Q, C, R, \{U_k\})$ are trained end-to-end via gradient methods (Adam). Differentiable SVDs circumvent instability in gradient computation through truncation steps [2606.01468].

The computational complexity per step is summarized as follows:

| Operation                | Standard Kalman          | CASSM                                 |
|--------------------------|-------------------------|----------------------------------------|
| Covariance storage       | $O(d^2)$                | $O(d \cdot m)$ (downdates, $m \ll d$) |
| Innovation inversion     | $O(n_k^3)$              | $O(m_k^3)$ ($m_k \ll n_k$)            |
| Overall time (per step)  | $O(d^3)$                | $O(d m^2)$ or $O(d m)$                |
| Overall memory           | $O(d^2)$                | $O(d m)$                              |

If $m$ and $r$ are small and kernel matvecs are $O(d \log d)$ (or $O(d)$), the total cost is $O(n_T d m)$ time and $O(d m)$ memory, linear in $d$ and favorable relative to cubic/quadratic baseline [2405.08971].

## 5. Model Selection, Learnable Projections, and Policies

In CASSM, projection matrices $U_k$ are not fixed but are treated as tunable policies:

- **Entropy-Based Policy:** Theoretically, the optimal $U_k$ for minimizing posterior entropy projects onto the top $M$ eigenvectors of the innovation covariance $S_k$ [2606.01468, Theorem 1]. However, direct computation is intractable for large $N$, so $U_k$ 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 $O(N)$. 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 $m$ [2405.08971].

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.,
  $$
  \sup_{f \in \mathcal{H}^\sigma \setminus \{0\}} \frac{|f(z) - \bar{y}(z)|}{\|f\|_{\mathcal{H}^\sigma}} = \sqrt{\sigma_C^2(z)},
  $$
  thus the posterior uncertainty envelopes both data and computational errors [2405.08971].
- **Empirical Results on Large-Scale Problems:**
    - On large-scale climate data ($d \approx 2.3 \times 10^5$, $n_T n_x \approx 4 \times 10^6$), CASSM achieves exponential improvement in MSE and NLL versus the number of CG steps $m$, 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 $\sim$0.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 [2606.01468].

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 $O(n_T d m)$ cost. Likelihood-based learning (e.g., EM) is possible, but gradient computation becomes involved when truncation is used [2405.08971].

## 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 ($N$) far exceeds the number of trials ($T_{trials}$), avoiding overfitting and under-calibration typical in overparameterized deep models [2606.01468].
- **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:**  
[2405.08971] "Computation-Aware Kalman Filtering and Smoothing"  
[2606.01468] "Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics"

Source: https://www.emergentmind.com/topics/computation-aware-state-space-model-cassm