Accepted to ProbNum 2026

Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics

CASSM brings scalable computation-aware Bayesian filtering to neural latent dynamics, including model selection and calibrated uncertainty in regimes with many neurons and few trials.

JR Huml / Jonathan Wenger / John P. Cunningham

Proceedings of the 2nd International Conference on Probabilistic Numerics, Lappeenranta, Finland, September 9-11, 2026

Paper arXiv Code ProbNum 2026
t = 0.00 sigma = 10 rho = 28 beta = 8/3
Abstract

Bayesian neural dynamics at modern recording scale

Bayesian dynamical latent variable models provide explicit priors and principled uncertainty, but classical formulations struggle as neural recordings grow in dimensionality. CASSM extends computation-aware Kalman filtering to learn and select models directly, using a tractable training objective designed for large state spaces.

The method targets the scale-imbalanced setting common in neuroscience, where the number of recorded neurons is much larger than the number of trials. In synthetic and real neural datasets, CASSM is competitive with deep latent dynamics models while improving uncertainty calibration over earlier scalable Bayesian approaches.

Method

Computation-aware filtering with learned model structure

Filtering and smoothing

PyTorch implementations of computation-aware filtering and smoothing for high-dimensional neural time series.

Model selection

A training loss and optimization scheme that learns latent dynamics, readout maps, and noise models rather than fixing hyperparameters by hand.

Uncertainty calibration

Computational uncertainty is carried through the inference pipeline so predictive accuracy and confidence can be evaluated together.

Open Source

Install CASSM and run the Lorenz tutorial

pip install cassm
Repository PyPI Tutorial
Package Python, MIT License
Core modules datasets, losses, metrics, models, utils
Example system Lorenz latent dynamics with neural readout
Venue

ProbNum 2026

ProbNum 2026 is the 2nd International Conference on Probabilistic Numerics, focused on statistically solving numerical problems and quantifying numerical error as computational uncertainty. The meeting takes place at LUT University in Lappeenranta, Finland.

Conference website
Citation

BibTeX

@inproceedings{huml2026cassm,
  title = {Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics},
  author = {Huml, JR and Wenger, Jonathan and Cunningham, John P.},
  booktitle = {Proceedings of the 2nd International Conference on Probabilistic Numerics},
  year = {2026},
  note = {ProbNum 2026},
  doi = {10.48550/arXiv.2606.01468},
  url = {https://arxiv.org/abs/2606.01468}
}