Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics

Technische Samenvatting: Computationeel Bewuste Kalman-filtering met Modelselectie voor Neurale Dynamica

Probleemstelling
Het artikel behandelt de computationele en statistische uitdagingen bij het modelleren van dynamische latente variabelen in moderne, grootschalige neurale opnames. Hoewel Bayesiaanse methoden (bijv. Gaussian Process Factor Analysis, GPFA) expliciete priors en onzekerheidskwantificering bieden, zijn hun kubische complexiteit ten opzichte van de temporele resolutie en hun kwadratische complexiteit ten opzichte van de ruimtelijke resolutie (aantal neuronen) onhaalbaar voor moderne datasets. Daartegenover staan overgeparameteriseerde deep learning-modellen (bijv. LFADS) die gunstig schalen, maar vaak een gebrek hebben aan robuuste onzekerheidskwantificering en een groot aantal trials vereisen om hun gebrek aan inductieve bias te overwinnen.

De auteurs identificeren een specifiek "schaal-imbalans regime" dat veel voorkomt in de neurowetenschappen (bijv. menselijke hersenstudies), waarbij het aantal opgenomen neuronen aanzienlijk groter is dan het aantal experimentele trials. In dit regime degraderen standaard deep learning-modellen door dataschaarste, terwijl traditionele Bayesiaanse benaderingen vaak niet in staat zijn om de fouten te compenseren die worden geïntroduceerd door hun eigen computationele beperkingen, wat leidt tot overmoedige ("variance-starved") onzekerheidsschattingen.

Methodologie: Computationeel Bewust State-Space Model (CASSM)
De auteurs stellen het Computation-Aware State-Space Model (CASSM) voor, een probabilistisch numeriek raamwerk dat de Computation-Aware Kalman Filter uitbreidt om modelselectie te bevatten. De kernmethodologie omvat:

1. Projectie en Computationele Onzekerheid: CASSM erkent dat de "ware" posterior computationeel onhaalbaar is om direct te berekenen. In plaats van de benaderingsfout te negeren, modelleert het expliciet "computationele onzekerheid". De methode projecteert hoogdimensionale observaties ($y_k \in \mathbb{R}^n$) naar een lager dimensionale ruimte ($\check{y}_k \in \mathbb{R}^{\check{n}}$) met behulp van een geleerde beleidsmatrix $S_k$. Dit vermindert de complexiteit van het inverteren van de innovatiematrix van $O(n^3)$ naar $O(\check{n}^3)$.
2. Low-Rank Covariantie-updates: Om geheugenbeperkingen aan te pakken ($O(d^2)$ voor state-covarianties), maakt CASSM gebruik van een low-rank downdate-structuur. Het slaat "skinny" low-rank vierkantswortels van de covariantie-updates op in plaats van de volledige matrices, waardoor de geheugeneisen worden verminderd tot $O(dr)$.
3. Modelselectie via PCA-analogie: De projectiebeleid $S_k$ wordt geïnterpreteerd via een informatietheoretisch oogpunt. Stelling 1 stelt vast dat de optimale actie (projectie) die de latente posterior-entropie minimaliseert, overeenkomt met de top-eigenvectoren van de innovatie-covariantie (analoog aan Principal Components Analysis). Het raamwerk leert deze acties om uit te lijnen met de richtingen van de grootste verwachte onzekerheid.
4. Nieuw Trainingsdoel: De auteurs introduceren een specifieke verliesfunctie afgeleid van de Evidence Lower Bound (ELBO). In tegen tegenstelling tot de standaard Negative Log-Likelihood (NLL), die kan degraderen onder projectie, balanceert de CASSM-objectief een data-fit term met een divergentie-term. Deze divergentie-term bestraft de afwijking van de benaderde posterior van de exacte filtering-distributie, waardoor onzekerheidsschattingen proportioneel breder worden naarmate de computationele benadering toeneemt.
5. Ruimtelijke Priors: Door Gaussian Process-representaties te vertalen naar State-Space Modellen, staat CASSM de expliciete codering van ruimtelijke priors (bijv. neuroncoördinaten of celtypen) toe via een ruimtelijke kernel $\Sigma^{(x)}$, een kenmerk dat slechts impliciet aanwezig is in eerdere methoden zoals GPFA of LFADS.
6. Preprocessing en Stabiliteit: De pipeline bevat variantie-stabiliserende transformaties (Anscombe) voor spike-data en maakt gebruik van differentiële singular value decompositie (gebaseerd op de Moore-Penrose pseudoinversie) om numerieke stabiliteit te garanderen tijdens de truncatie van low-rank updates.

Belangrijkste Bijdragen

• CASSM-raamwerk: Een schaalbare, probabilistische numerieke methode die specifbest is ontworpen voor het schaal-imbalans regime (trials $\ll$ neuronen).
• Tractable Inferentie met Onzekerheidskalibratie: De methode levert tractabele inferentie in grote state-spaces terwijl het gekalibreerde onzekerheid biedt die rekening houdt met computationele benaderingsfouten, waardoor de overmoedigheid die typisch is voor andere benaderingen wordt vermeden.
• Nieuwe Verliesfunctie: Een trainingsverlies dat efficiënte hyperparameteroptimalisatie en modelselectie mogelijk maakt door divergentie van de exacte filtering-distributie te bestraffen, waardoor het model informatieve data-projecties kan leren.
• Expliciete Ruimtelijke Modellering: Een techniek om expliciet ruimtelijke priors over neuronen in een state-space model te integreren, wat CASSM onderscheidt van latente variabele modellen die ruimtelijke structuur impliciet behandelen.

Resultaten
De auteurs evalueren CASSM op zowel synthetische als real-world datasets:

• Synthetische Lorenz Data: CASSM vertoont lineaire schaling ten opzichte van het aantal neuronen en presteert beter dan GPFA en bGPFA in computationele tijd bij grote schalen ($10^4$ neuronen). In het schaal-imbalans regime behoudt CASSM een concurrerende voorspellende prestatie (MSE) vergeleken met LFADS, terwijl LFADS snel degradeert wanneer trials beperkt zijn.
• Onzekerheidskalibratie: Experimenten met dekking tonen aan dat CASSM significant beter gekalibreerde 95% betrouwbaarheidsintervallen biedt vergeleken met bGPFA en GPFA, die de neiging hebben tot overmoedigheid.
• Echte Neurale Data (Neural Latents Benchmark): Op datasets zoals Area 2 Bump en DMFC presteert CASSM competitief met LFADS en bGPFA. In de "MC Maze" dataset (140 neuronen, 1720 trials) bereikt CASSM vergelijkbare co-smoothing scores als LFADS.
• Zebravis-dataset: Op een grootschalige dataset met ~88.000 neuronen en slechts 96 trials faalt bGPFA door out-of-memory fouten. CASSM verwerkt de data succesvol, en wanneer het ruimtelijke priors gebruikt (CASSM-SP), bereikt het een prestatie die statistisch vergelijkbaar is met LFADS, terwijl het gekalibreerde onzekerheid biedt.

Significantie en Claims
Het artikel beweert dat CASSM een "routekaart" biedt voor onderzoekers om te kiezen tussen dynamische latente variabele modellen op basis van de eigenschappen van de dataset. De auteurs beargumenteren dat terwijl deep learning-methoden (LFADS) superieur zijn wanneer trials het aantal neuronen overtreffen, Bayesiaanse methoden essentieel worden in het schaal-imbalans regime dat typerend is voor grootschalige neurowetenschappen. CASSM overbrugt deze kloof door een methode te bieden die computationeel schaalbaar is zoals deep networks, maar de inductieve biases en onzekerheidskwantificering van Bayesiaanse inferentie behoudt. De auteurs merken bescheiden op dat CASSM zeer gespecialiseerd is voor dit specifieke regime en een lineaire methode is, wat betekent dat het moeite kan hebben als de onderliggende dynamica zeer niet-lineair is of als de rank van de benadering slecht is gekozen. Echter, in het doelregime biedt het een competitief alternatief voor data-hongerige deep networks met superieure onzekerheidskalibratie.