📊 statistics

Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics

This paper introduces the Computation-Aware State-Space Model (CASSM), a framework that extends Bayesian latent variable modeling to large-scale neural datasets by incorporating computational uncertainty and model selection, thereby achieving competitive performance with deep networks while offering superior uncertainty calibration in data-scarce regimes.

Original authors: JR Huml, Jonathan Wenger, John P. Cunningham

Published 2026-06-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: JR Huml, Jonathan Wenger, John P. Cunningham

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to understand the weather patterns of a massive, bustling city. You have two main tools to help you:

The Super-Computer Weather Model: This is a massive, complex simulation that needs a huge amount of historical data to work well. If you give it thousands of years of weather records, it predicts the future perfectly. But if you only give it data from last week, it gets confused and starts making wild guesses. It's also very expensive to run.
The Experienced Local Forecaster: This person uses simple rules of thumb (inductive bias) and knows how the city usually behaves. They don't need thousands of years of data; they can make a decent prediction with just a few days of records. However, they aren't as precise as the super-computer when there is lots of data.

The Problem:
In neuroscience, scientists are trying to do the same thing: predict how a brain works based on the "spikes" (electrical signals) of neurons.

The Old Way (Deep Learning/Neural Nets): Like the Super-Computer. It's great when you have millions of trials (repetitions of an experiment) but fails when you have few trials.
The Traditional Bayesian Way: Like the Local Forecaster. It's great at handling uncertainty and small datasets, but it's computationally slow. As the number of neurons grows (from hundreds to thousands), the math becomes so heavy it crashes the computer.

The New Solution: CASSM
The authors introduce a new method called CASSM (Computation-Aware State-Space Model). Think of CASSM as a "Smart Forecaster with a Reality Check."

Here is how it works, using simple analogies:

1. The "Reality Check" (Computational Uncertainty)

Most methods try to simplify complex math to make it run faster. The problem is, they often forget to tell you how much they simplified it. It's like a GPS that says, "Turn left in 500 feet," but doesn't tell you it's using a 10-year-old map. You might drive into a construction zone.

CASSM is different. It explicitly calculates the "error" caused by simplifying the math.

The Metaphor: Imagine you are estimating the weight of a giant elephant.
- Old Method: You guess it weighs 5 tons and say, "I'm 100% sure." (Overconfident).
- CASSM: You guess it weighs 5 tons, but you add a note: "Because I had to use a small scale, my estimate might be off by 2 tons."
- Why it matters: As the computer works harder to be more accurate, CASSM shrinks that "off by" note. If the computer is lazy, the note gets bigger. This prevents scientists from trusting a bad guess.

2. The "Scale-Imbalanced" Regime

The paper focuses on a specific, difficult situation: When you have a huge number of neurons (the city) but very few trials (the weather data).

In the human brain, you might have billions of neurons, but you can't run billions of experiments on a single person.
Deep learning models (like LFADS) usually fail here because they are "data-hungry." They starve without enough trials.
CASSM is designed specifically for this starvation diet. It uses "priors" (smart assumptions about how neurons behave) to fill in the gaps where data is missing.

3. The "Smart Lens" (Model Selection)

One of CASSM's coolest features is how it decides what to look at.

The Analogy: Imagine you are looking at a crowded room through a small keyhole. You can only see a few people at a time.
Old Methods: They might look through the keyhole randomly.
CASSM: It learns to move the keyhole to the most interesting parts of the room. It automatically figures out which neurons are most important to track to understand the whole system. It does this by trying to minimize "confusion" (entropy) in its predictions.

4. The "Spatial Map"

CASSM also lets scientists plug in a map of where the neurons are located in the brain.

The Analogy: If you are predicting traffic, it helps to know that a road is a highway vs. a side street.
The Innovation: CASSM can use the physical location of neurons (or their cell type) as a hint for how they should behave. This is something older methods did only implicitly or not at all.

What Did They Find?

The authors tested CASSM on both fake data (simulated neurons) and real data (monkeys reaching for objects, zebrafish brains).

When data is plentiful: Deep learning models (like LFADS) are still the kings of prediction. CASSM is good, but not the absolute best.
When data is scarce (The "Scale-Imbalanced" zone): This is where CASSM shines.
- It predicts almost as well as the deep learning models.
- It is much faster than traditional Bayesian methods (like GPFA) when the number of neurons gets huge (thousands).
- Crucially: It gives you a "confidence score" that is actually honest. Other methods often pretend to be sure when they are actually guessing; CASSM admits when it's unsure.

The Bottom Line

CASSM is a new tool for neuroscientists that bridges the gap between "fast but dumb" deep learning and "smart but slow" traditional statistics. It is specifically built for the future of neuroscience, where we will record from thousands of neurons but can't run thousands of experiments. It allows researchers to get reliable answers with honest confidence levels, even when the math is too heavy to solve perfectly.

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

Problem Statement
The paper addresses the computational and statistical challenges of modeling dynamical latent variables in modern, large-scale neural recordings. While Bayesian methods (e.g., Gaussian Process Factor Analysis, GPFA) offer explicit priors and uncertainty quantification, their cubic complexity with respect to temporal resolution and quadratic complexity with respect to spatial resolution (number of neurons) render them infeasible for modern datasets. Conversely, overparameterized deep learning models (e.g., LFADS) scale favorably but often lack robust uncertainty quantification and require large numbers of trials to overcome their lack of inductive bias.

The authors identify a specific "scale-imbalanced regime" prevalent in neuroscience (e.g., human brain studies), where the number of recorded neurons significantly exceeds the number of experimental trials. In this regime, standard deep learning models degrade due to data scarcity, while traditional Bayesian approximations often fail to account for the errors introduced by their own computational constraints, leading to overconfident ("variance-starved") uncertainty estimates.

Methodology: Computation-Aware State-Space Model (CASSM)
The authors propose the Computation-Aware State-Space Model (CASSM), a probabilistic numerical framework that extends the Computation-Aware Kalman Filter (CAKF) to include model selection. The core methodology involves:

Projection and Computational Uncertainty: CASSM acknowledges that the "true" posterior is computationally infeasible to compute directly. Instead of ignoring the approximation error, it explicitly models "computational uncertainty." The method projects high-dimensional observations ( $y_k \in \mathbb{R}^n$ ) into a lower-dimensional space ( $\check{y}_k \in \mathbb{R}^{\check{n}}$ ) using a learned policy matrix $S_k$ . This reduces the complexity of inverting the innovation matrix from $O(n^3)$ to $O(\check{n}^3)$ .
Low-Rank Covariance Updates: To address memory constraints ( $O(d^2)$ for state covariances), CASSM utilizes a low-rank downdate structure. It stores "skinny" low-rank square roots of the covariance updates rather than the full matrices, reducing memory requirements to $O(dr)$.
Model Selection via PCA Analogue: The projection policy $S_k$ is interpreted through an information-theoretic lens. Theorem 1 establishes that the optimal action (projection) that minimizes latent posterior entropy corresponds to the top eigenvectors of the innovation covariance (analogous to Principal Components Analysis). The framework learns these actions to align with the directions of greatest expected uncertainty.
Novel Training Objective: The authors introduce a specific loss function derived from the Evidence Lower Bound (ELBO). Unlike standard Negative Log-Likelihood (NLL) which may degrade under projection, the CASSM objective balances a data-fit term with a divergence term. This divergence term penalizes the deviation of the approximate posterior from the exact filtering distribution, ensuring that uncertainty estimates widen appropriately as computational approximation increases.
Spatial Priors: By translating Gaussian Process representations to State-Space Models, CASSM allows for the explicit encoding of spatial priors (e.g., neuron coordinates or cell types) via a spatial kernel $\Sigma^{(x)}$ , a feature only implicit in previous methods like GPFA or LFADS.
Preprocessing and Stability: The pipeline includes variance-stabilizing transforms (Anscombe) for spike data and utilizes differentiable singular value decomposition (based on the Moore-Penrose pseudoinverse) to ensure numerical stability during the truncation of low-rank updates.

Key Contributions

CASSM Framework: A scalable, probabilistic numerical method designed specifically for the scale-imbalanced regime (trials $\ll$ neurons).
Tractable Inference with Uncertainty Calibration: The method yields tractable inference in large state-spaces while providing calibrated uncertainty that accounts for computational approximation errors, avoiding the overconfidence typical of other approximations.
Novel Loss Function: A training loss that enables efficient hyperparameter optimization and model selection by penalizing divergence from the exact filtering distribution, allowing the model to learn informative data projections.
Explicit Spatial Modeling: A technique to explicitly incorporate spatial priors over neurons within a state-space model, differentiating it from latent variable models that treat spatial structure implicitly.

Results
The authors evaluate CASSM on both synthetic and real-world datasets:

Synthetic Lorenz Data: CASSM demonstrates linear scaling with respect to the number of neurons, outperforming GPFA and bGPFA in computational time at large scales ( $10^4$ neurons). In the scale-imbalanced regime, CASSM maintains competitive predictive performance (MSE) compared to LFADS, whereas LFADS degrades rapidly when trials are limited.
Uncertainty Calibration: Coverage experiments show that CASSM provides significantly better-calibrated 95% confidence intervals compared to bGPFA and GPFA, which tend to be overconfident.
Real Neural Data (Neural Latents Benchmark): On datasets like Area 2 Bump and DMFC, CASSM performs competitively with LFADS and bGPFA. In the "MC Maze" dataset (140 neurons, 1720 trials), CASSM achieves comparable co-smoothing scores to LFADS.
Zebrafish Dataset: On a large-scale dataset with ~88,000 neurons and only 96 trials, bGPFA fails due to out-of-memory errors. CASSM successfully processes the data, and when utilizing spatial priors (CASSM-SP), it achieves performance statistically similar to LFADS while providing calibrated uncertainty.

Significance and Claims
The paper claims that CASSM provides a "roadmap" for researchers to choose dynamical latent variable models based on dataset properties. The authors argue that while deep learning methods (LFADS) are superior when trials outnumber neurons, Bayesian methods become essential in the scale-imbalanced regime typical of large-scale neuroscience. CASSM bridges this gap by offering a method that is computationally scalable like deep networks but retains the inductive biases and uncertainty quantification of Bayesian inference. The authors modestly note that CASSM is highly specialized for this specific regime and is a linear method, meaning it may struggle if the underlying dynamics are highly nonlinear or if the rank of the approximation is chosen poorly. However, in the target regime, it offers a competitive alternative to data-hungry deep networks with superior uncertainty calibration.