Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics

2026-05-31

AI Summary

This work addresses the challenges faced by conventional Bayesian approaches in neural dynamics modeling under high-dimensional, small-sample regimes—where the number of experimental trials is far fewer than the number of neurons—namely poor computational efficiency, miscalibrated uncertainty estimates, and difficulties in model selection. To overcome these limitations, the authors propose a Computationally Aware State Space Model (CASSM), which introduces computational awareness into the state space model selection process for the first time. By integrating Kalman filtering with explicit model selection, modeling of computational uncertainty, and a novel loss function, CASSM achieves efficient and scalable inference while preserving the advantages of Bayesian modeling. Experiments demonstrate that the method attains prediction performance comparable to deep neural networks on both synthetic and real neural data, while substantially improving the calibration of predictive uncertainties.

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.

Problem

• neural dynamics
• Bayesian methods
• model selection
• computational uncertainty
• scale-imbalanced regime

Innovation

• Computation-Aware Kalman Filtering
• Model Selection
• Computational Uncertainty
• State-Space Model
• Neural Dynamics