Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics

Sintesi Tecnica: Filtraggio di Kalman Consapevole della Computazione con Selezione del Modello per la Dinamica Neurale

Definizione del Problema
Il paper affronta le sfide computazionali e statistiche della modellazione di variabili latenti dinamiche nelle moderne registrazioni neurali su larga scala. Sebbene i metodi bayesiani (ad esempio, Gaussian Process Factor Analysis, GPFA) offrano prior espliciti e quantificazione dell'incertezza, la loro complessità cubica rispetto alla risoluzione temporale e la complessità quadratica rispetto alla risoluzione spaziale (numero di neuroni) li rendono impraticabili per i moderni dataset. Al contrario, i modelli di deep learning sovraparametrizzati (ad esempio, LFADS) scalano favorevolmente ma spesso mancano di una robusta quantificazione dell'incertezza e richiedono un gran numero di trial per superare la loro mancanza di bias induttivo.

Gli autori identificano un regime specifico di "squilibrio di scala" prevalente nelle neuroscienze (ad esempio, studi sul cervello umano), dove il numero di neuroni registrati supera significativamente il numero di trial sperimentali. In questo regime, i modelli standard di deep learning degradano a causa della scarsità di dati, mentre le tradizionali approssimazioni bayesiane spesso non riescono a tenere conto degli errori introdotti dalle proprie restrizioni computazionali, portando a stime di incertezza eccessivamente sicure ("prive di varianza").

Metodologia: Modello di Stato-Spazio Consapevole della Computazione (CASSM)
Gli autori propongono il Modello di Stato-Spazio Consapevole della Computazione (CASSM), un framework numerico probabilistico che estende il Filtro di Kalman Consapevole della Computazione (CAKF) per includere la selezione del modello. La metodologia centrale prevede:

1. Proiezione e Incertezza Computazionale: Il CASSM riconosce che il posteriore "vero" è computazionalmente impraticabile da calcolare direttamente. Invece di ignorare l'errore di approssimazione, esso modella esplicitamente l' "incertezza computazionale". Il metodo proietta osservazioni ad alta dimensione ($y_k \in \mathbb{R}^n$) in uno spazio a dimensione ridotta ($\check{y}_k \in \mathbb{R}^{\check{n}}$) utilizzando una matrice di policy appresa $S_k$. Ciò riduce la complessità dell'inversione della matrice di innovazione da $O(n^3)$ a $O(\check{n}^3)$.
2. Aggiornamenti di Covarianza a Basso Rango: Per affrontare i vincoli di memoria ($O(d^2)$ per le covarianze di stato), il CASSM utilizza una struttura di downdate a basso rango. Memorizza radici quadrate "snelle" a basso rango degli aggiornamenti di covarianza invece delle matrici complete, riducendo i requisiti di memoria a $O(dr)$.
3. Selezione del Modello tramite Analogo PCA: La policy di proiezione $S_k$ è interpretata attraverso una lente informazionale. Il Teorema 1 stabilisce che l'azione ottimale (proiezione) che minimizza l'entropia del posteriore latente corrisponde agli autovettori superiori della covarianza di innovazione (analogo alla Principal Component Analysis). Il framework apprende queste azioni per allinearsi alle direzioni di maggiore incertezza attesa.
4. Nuova Funzione di Obiettivo: Gli autori introducono una specifica funzione di perdita derivata dall'Evidence Lower Bound (ELBO). A differenza della standard Negative Log-Likelihood (NLL), che può degradare sotto proiezione, l'obiettivo CASSM bilancia un termine di adattamento ai dati con un termine di divergenza. Questo termine di divergenza penalizza la deviazione del posteriore approssimato dalla distribuzione di filtraggio esatta, garantendo che le stime di incertezza si allarghino adeguatamente all'aumentare dell'approssimazione computazionale.
5. Prior Spaziali: Traducendo le rappresentazioni dei Processi Gaussiani in Modelli di Stato-Spazio, il CASSM consente la codifica esplicita di prior spaziali (ad esempio, coordinate dei neuroni o tipi cellulari) tramite un kernel spaziale $\Sigma^{(x)}$, una caratteristica che è solo implicita nei metodi precedenti come GPFA o LFADS.
6. Pre-elaborazione e Stabilità: La pipeline include trasformazioni stabilizzanti della varianza (Anscombe) per i dati di spike e utilizza la decomposizione a valori singolari differenziabile (basata sulla pseudoinversa di Moore-Penrose) per garantire la stabilità numerica durante la troncatura degli aggiornamenti a basso rango.

Contributi Chiave

• Framework CASSM: Un metodo numerico probabilistico scalabile progettato specificamente per il regime di squilibrio di scala (trial $\ll$ neuroni).
• Inferenza Trattabile con Calibrazione dell'Incertezza: Il metodo fornisce inferenza trattabile in spazi di stato elevati fornendo al contempo un'incertezza calibrata che tiene conto degli errori di approssimazione computazionale, evitando l'eccessiva sicurezza tipica di altre approssimazioni.
• Nuova Funzione di Perdita: Una funzione di perdita che consente l'ottimizzazione efficiente degli iperparametri e la selezione del modello penalizzando la divergenza dalla distribuzione di filtraggio esatta, permettendo al modello di apprendere proiezioni dei dati informative.
• Modellazione Spaziale Esplicita: Una tecnica per incorporare esplicitamente prior spaziali sui neuroni all'interno di un modello di stato-spazio, differenziandolo dai modelli di variabili latenti che trattano la struttura spaziale implicitamente.

Risultati
Gli autori valutano il CASSM sia su dati sintetici che su dataset del mondo reale:

• Dati Lorenz Sintetici: Il CASSM dimostra una scalabilità lineare rispetto al numero di neuroni, superando GPFA e bGPFA in tempo computazionale a grandi scale ($10^4$ neuroni). Nel regime di squilibrio di scala, il CASSM mantiene prestazioni predittive competitive (MSE) rispetto a LFADS, mentre LFADS degrada rapidamente quando i trial sono limitati.
• Calibrazione dell'Incertezza: Gli esperimenti di copertura mostrano che il CASSM fornisce intervalli di confidenza al 95% significativamente meglio calibrati rispetto a bGPFA e GPFA, che tendono a essere eccessivamente sicuri.
• Dati Neurali Reali (Neural Latents Benchmark): Su dataset come Area 2 Bump e DMFC, il CASSM performa competitivamente con LFADS e bGPFA. Nel dataset "MC Maze" (140 neuroni, 1720 trial), il CASSM raggiunge punteggi di co-smoothing comparabili a quelli di LFADS.
• Dataset Zebrafish: Su un dataset su larga scala con circa 88.000 neuroni e solo 96 trial, bGPFA fallisce a causa di errori di out-of-memory. Il CASSM elabora con successo i dati e, utilizzando i prior spaziali (CASSM-SP), ottiene prestazioni statisticamente simili a LFADS fornendo al contempo un'incertezza calibrata.

Significatività e Rivendicazioni
Il paper sostiene che il CASSM fornisca una "tabella di marcia" per i ricercatori per scegliere i modelli di variabili latenti dinamiche in base alle proprietà del dataset. Gli autori argomentano che, mentre i metodi di deep learning (LFADS) sono superiori quando i trial superano i neuroni, i metodi bayesiani diventano essenziali nel regime di squilibrio di scala tipico delle neuroscienze su larga scala. Il CASSM colma questo divario offrendo un metodo che è computazionalmente scalabile come le reti profonde, ma che mantiene i bias induttivi e la quantificazione dell'incertezza dell'inferenza bayesiana. Gli autori notano con modestia che il CASSM è altamente specializzato per questo specifico regime ed è un metodo lineare, il che significa che potrebbe avere difficoltà se la dinamica sottostante è altamente non lineare o se il rango dell'approssimazione è scelto male. Tuttavia, nel regime target, offre un'alternativa competitiva ai network di deep learning affamati di dati con una superiore calibrazione dell'incertezza.