Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics
This work addresses the need for scalable Bayesian dynamical latent variable models in neuroscience that can handle modern-sized datasets without sacrificing uncertainty quantification.
The authors propose a computation-aware Kalman filtering framework (CASSM) for neural dynamics that handles model selection and scales to large state-spaces, achieving competitive predictive performance with deep networks while providing better uncertainty calibration in the scale-imbalanced regime where trials are fewer than neurons.
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.