Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics

Abstract: Due to their explicit priors and ability to model uncertainty, Bayesian methods have played a major role in dynamical latent variable modeling of single-cell neural recordings. However, modern-sized datasets have made overparameterized deep networks the preferred methods of choice due to their predictive power and favorable computational scaling. While many posterior approximations exist, all incur approximation errors. Recent work accounts for this error in the form of computational uncertainty but comes at the cost of quadratic complexity and assumes fixed model hyperparameters. Here we extend this development to model selection, including a novel training loss and optimization scheme, which yields tractable inference in large state-spaces. We introduce a framework, the Computation-Aware State-Space Model (CASSM), specifically designed for the scale-imbalanced regime, where the number of trials is significantly lower than the number of recorded neurons. In this regime, for both synthetic and real data, we show that our method is competitive with data-hungry deep networks, with significantly improved uncertainty calibration over previous attempts to scale Bayesian methods. Our experiments provide a roadmap to neuroscience researchers in choosing from a host of potential dynamical latent variable models given key dataset properties and constraints.

• The paper introduces CASSM, a novel state-space model that optimizes computational cost against approximation-induced uncertainty in neural dynamics inference.
• It leverages matrix-free, block-sparse operator representations and explicit spatial priors to scale efficiently to ultra-high neuron counts in trial-limited settings.
• CASSM achieves robust predictive performance and principled model selection through a variational loss that aligns approximate posteriors with true filtering distributions.

Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics: A Technical Overview

Introduction and Problem Motivation

Modern neuroscience increasingly relies on high-density multi-electrode arrays and whole-brain imaging to capture large-scale, high-dimensional population activity. The key challenge is to extract interpretable low-dimensional latent neural dynamics from these massive, often trial-limited datasets. Traditional Bayesian state-space models like Gaussian Process Factor Analysis (GPFA) provide strong inductive biases and robust uncertainty quantification, but their cubic computational complexity in time and quadratic complexity in neuron count limit scalability. In contrast, deep neural latent models such as LFADS offer favorable scalability and superior predictive performance in data-rich regimes but often lack meaningful uncertainty estimates and require many independent trials to offset their statistical inefficiency.

The paper introduces the Computation-Aware State-Space Model (CASSM), which generalizes the Computation-Aware Kalman Filter to enable tractable inference, uncertainty quantification, and automated model selection in the "scale-imbalanced" regime—where the number of trials is much smaller than the number of recorded neurons—by explicitly modeling the uncertainty induced by approximate inference.

Model Formulation and Theoretical Insights

CASSM frames the problem of inference and model selection in latent linear dynamical systems as an optimization over operator projections that trade computational cost for approximation-induced uncertainty. The methodology extends the notion of "computational uncertainty" introduced in prior computation-aware filtering frameworks, but crucially allows for data-driven model selection via a novel variational loss.

Model selection is achieved by optimizing a data-fit term together with a divergence penalty between the computationally tractable approximate posterior and the (infeasible) true filtering distribution. The optimal approximate projection directions are shown to coincide with the principal components (top eigenvectors) of the innovation covariance, establishing a formal connection to PCA in the context of sequential filtering. This guarantees that the learned projection maximizes uncertainty reduction at every time step, subject to computational constraints.

Figure 1: The Grassmann distance drops as the learned projection aligns with the principal subspace of the innovation covariance, confirming convergence to the entropy-minimizing direction set.

CASSM uses a low-rank compressed representation of latent and observation covariances and leverages kernel-parameterized priors to exploit spatial structure. This enables explicit encoding of spatial priors—such as cell type or anatomical coordinates—which prior approaches enforced only implicitly.

Algorithmic Implementation

CASSM relies on matrix-free, block-sparse operator representations to ensure memory and compute efficiency for large state spaces. The training objective—a collapsed ELBO—enforces fidelity to the true-data log-likelihood while regulating posterior collapse via a KL penalty between the approximate filtered posterior and its full-rank counterpart. This objective admits efficient evaluation and differentiation, and the use of differentiable SVD-based pseudoinverses stabilizes low-rank truncation during backpropagation.

The approach is robust to degenerate or ill-conditioned singular values during training, in contrast to previous low-rank Kalman approaches, which are numerically unstable under end-to-end gradient optimization over the projection basis.

Experimental Evaluation

Synthetic Dynamical Systems

On benchmark synthetic datasets (e.g., Lorenz system with non-linear latent dynamics), CASSM achieves strong test MSE and calibrated uncertainty coverage in the limiting regime of low trial counts and high neuron numbers. In the low-trial regime, Bayesian methods such as CASSM display marked robustness in predictive accuracy relative to data-hungry deep methods.

Figure 2: At large scales, CASSM outpaces GPFA and Kalman Filter by almost a day, and runs several hours faster than scalable Bayesian competitors, entering linear scaling around $10^4$ neurons, confirming practical tractability.

Figure 3: Bayesian priors in CASSM and bGPFA yield strong inductive bias gains, maintaining predictive performance as trials-per-neuron decrease; LFADS degrades rapidly in this regime.

Figure 4: In the data-abundant regime ($10^2$ neurons, trials = neurons), LFADS outperforms all Bayesian contenders; when uncertainty quantification is unnecessary, CASSM's utility is limited.

Uncertainty Quantification

CASSM explicitly models the increase in uncertainty due to both the projection (approximation) and intrinsic measurement noise, allowing principled separation of credible intervals due to computation limits from those intrinsic to the data.

Figure 5: Uncertainty bands widen as projection dimensionality drops from 30 to 2, visualizing computational uncertainty induced by approximation.

Figure 6: CASSM's posterior estimates show progressively increasing uncertainty under coarser subspace approximations, supporting faithful uncertainty calibration.

Real Neural Data

Experiments on Neural Latents Benchmark recordings from primate cortex and on massive light-sheet zebrafish calcium imaging data confirm the practical value of CASSM across real-world settings.

• In regimes where the number of trials is on par with or exceeds neuron count, non-linear deep latent models like LFADS retain a significant advantage in both held-out neuron prediction and behavioral decoding metrics.
• In the scale-imbalanced regime (tens of thousands of neurons, $\ll$100 trials), CASSM saturates or exceeds the computational scaling and predictive accuracy of deep models, while also maintaining credible uncertainty intervals.

  Figure 7: CASSM-reconstructed hand velocities, across four bump trial directions, exhibit realistic kinematic profiles with uncertainty envelopes, demonstrating behavioral decoding relevance.

Notably, inclusion of explicit spatial priors in CASSM on the zebrafish data narrows the performance gap to LFADS, at a similar wall-clock cost, while bGPFA fails due to out-of-memory issues for such ultra-high neuron counts.

Tradeoff Analysis and Model Selection Guidance

A key contribution is operational guidance in model selection across latent variable models. A decision analysis is synthesized based on neural scale, data/trial richness, and need for calibrated uncertainty:

Figure 8: Decision analysis tree maps neural data regimes to latent inference model selection, leveraging empirical findings from synthetic and real data.

• For low neuron, many-trial regimes, nonlinear VAEs (LFADS) are preferred when uncertainty is less essential.
• For scale-imbalanced or spatially structured data where computational cost or uncertainty accountability is critical, computation-aware methods like CASSM dominate.

Theoretical Implications and Future Directions

CASSM establishes rigorous variational bounds on approximation-induced uncertainty and connects dynamical inference to information-theoretic optimality principles (entropy minimization in time). The theoretical framework clarifies tradeoffs between projection-induced bias and computational feasibility. Explicit kernelized spatial priors create new opportunities for structured group comparisons in high-dimensional population analyses, and the matrix-free, operator-centric implementation blueprint can generalize to other spatiotemporal modalities.

Anticipated future work includes extending CASSM to flexible likelihoods, non-linear dynamical systems, and further reducing memory overhead via advances in randomized linear algebra. Applications to climate, sensor networks, and non-neural high-dimensional time series are envisaged.

Conclusion

CASSM supplies a computation-aware, uncertainty-calibrated latent dynamical modeling approach geared for large-scale, trial-limited neural datasets. Its integration of model selection with computational uncertainty quantification not only provides robust, scalable inference but also supports principled choices among competing latent variable models in contemporary neuroscience analysis. The proposed formalism and empirical benchmarking suggest practical applicability well beyond neuroscience, with theory-driven extensions toward nonlinear and non-Gaussian latent processes remaining an important avenue for future advancement.