Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics

Resumen Técnico: Filtrado de Kalman con Conciencia Computacional y Selección de Modelos para la Dinámica Neuronal

Planteamiento del Problema
El artículo aborda los desafíos computacionales y estadísticos de modelar variables latentes dinámicas en registros neuronales modernos y de gran escala. Mientras que los métodos bayesianos (por ejemplo, el Análisis de Factores de Procesos Gaussianos, GPFA) ofrecen priors explícitos y cuantificación de la incertidumbre, su complejidad cúbica con respecto a la resolución temporal y su complejidad cuadrática con respecto a la resolución espacial (número de neuronas) los hacen inviables para conjuntos de datos modernos. Por el contrario, los modelos de aprendizaje profundo sobreparametrizados (por ejemplo, LFADS) escalan favorablemente, pero a menudo carecen de una cuantificación robusta de la incertidumbre y requieren un gran número de ensayos para superar su falta de sesgo inductivo.

Los autores identifican un "régimen de desequilibrio de escala" prevalente en la neurociencia (por ejemplo, estudios del cerebro humano), donde el número de neuronas registradas supera significamente al número de ensayos experimentales. En este régimen, los modelos de aprendizaje profundo estándar se degradan debido a la escasez de datos, mientras que las aproximaciones bayesianas tradicionales a menudo no logran contabilizar los errores introducidos por sus propias restricciones computacionales, lo que conduce a estimaciones de incertidumbre sobreconfiadas ("carentes de varianza").

Metodología: Modelo de Espacio de Estados con Conciencia Computacional (CASSM)
Los autores proponen el Modelo de Espacio de Estados con Conciencia Computacional (CASSM), un marco numérico probabilístico que extiende el Filtro de Kalman con Conciencia Computacional (CAKF) para incluir la selección de modelos. La metodología central involucra:

1. Proyección e Incertidumbre Computacional: El CASSM reconoce que el posterior "verdadero" es computacionalmente inviable de calcular directamente. En lugar de ignorar el error de aproximación, modela explícitamente la "incertidumbre computacional". El método proyecta observaciones de alta dimensión ($y_k \in \mathbb{R}^n$) en un espacio de menor dimensión ($\check{y}_k \in \mathbb{R}^{\check{n}}$) utilizando una matriz de política aprendida $S_k$. Esto reduce la complejidad de invertir la matriz de innovación de $O(n^3)$ a $O(\check{n}^3)$.
2. Actualizaciones de Covarianza de Bajo Rango: Para abordar las restricciones de memoria ($O(d^2)$ para las covarianzas de estado), el CASSM utiliza una estructura de actualización de bajo rango (low-rank downdate). Almacena raíces cuadradas de bajo rango "delgadas" de las actualizaciones de la covarianza en lugar de las matrices completas, reduciendo los requisitos de memoria a $O(dr)$.
3. Selección de Modelos vía Analogía de PCA: La política de proyección $S_k$ se interpreta a través de una lente de información teórica. El Teorema 1 establece que la acción óptima (proyección) que minimiza la entropía del posterior latente corresponde a los autovectores superiores de la covarianza de la innovación (análogo al Análisis de Componentes Principales o PCA). El marco aprende estas acciones para alinearse con las direcciones de mayor incertidumbre esperada.
4. Objetivo de Entrenamiento Novedoso: Los autores introducen una función de pérdida específica derivada de la Cota Inferior de la Evidencia (ELBO). A diferencia de la Verosimilitud Negativa de Logaritmo (NLL) estándar, que puede degradarse bajo la proyección, el objetivo de CASSM equilibra un término de ajuste de datos con un término de divergencia. Este término de divergencia penaliza la desviación del posterior aproximado de la distribución de filtrado exacta, asegurando que las estimaciones de incertidumbre se amplíen apropiadamente a medida que aumenta la aproximación computacional.
5. Priors Espaciales: Al traducir las representaciones de Procesos Gaussianos a Modelos de Espacio de Estados, el CASSM permite la codificación explícita de priors espaciales (por ejemplo, coordenadas neuronales o tipos celulares) mediante un kernel espacial $\Sigma^{(x)}$, una característica que es solo implícita en métodos anteriores como GPFA o LFADS.
6. Preprocesamiento y Estabilidad: El flujo de trabajo incluye transformaciones de estabilización de la varianza (Anscombe) para datos de picos (spikes) y utiliza la descomposición en valores singulares diferenciable (basada en la pseudoinversa de Moore-Penrose) para asegurar la estabilidad numérica durante la truncación de las actualizaciones de bajo rango.

Contribuciones Clave

• Marco CASSM: Un método numérico escalable y probabilístico diseñado específicamente para el régimen de desequilibrio de escala (ensayos $\ll$ neuronas).
• Inferencia Tractable con Calibración de Incertidumbre: El método produce inferencia tratable en espacios de estados de gran tamaño proporcionando, además, una incertidumbre calibrada que tiene en cuenta los errores de aproximación computacional, evitando la sobreconfianza típica de otras aproximaciones.
• Función de Pérdida Novedosa: Una función de pérdida de entrenamiento que permite la optimización eficiente de hiperparámetros y la selección de modelos al penalizar la divergencia de la distribución de filtrado exacta, permitiendo que el modelo aprenda proyecciones de datos informativas.
• Modelado Espacial Explícito: Una técnica para incorporar explícitamente priors espaciales sobre las neuronas dentro de un modelo de espacio de estados, diferenciándose de los modelos de variables latentes que tratan la estructura espacial de forma implícita.

Resultos
Los autores evalúan el CASSM tanto en datos sintéticos como en datos reales:

• Datos de Lorenz Sintéticos: El CASSM demuestra un escalado lineal con respecto al número de neuronas, superando a GPFA y bGPFA en tiempo computacional en escalas grandes ($10^4$ neuronas). En el régimen de desequilibrio de escala, el CASSM mantiene un rendimiento predictivo competitivo (MSE) comparado con LFADS, mientras que LFADS se degrada rápidamente cuando los ensayos son limitados.
• Calibración de la Incertidumbre: Los experimentos de cobertura muestran que el CASSM proporciona intervalos de confianza del 95% significativamente mejor calibrados en comparación con bGPFA y GPFA, que tienden a ser sobreconfiados.
• Datos Neuronales Reales (Benchmark de Latentes Neurales): En conjuntos de datos como Area 2 Bump y DMFC, el CASSM rinde de manera competitiva con LFADS y bGPFA. En el conjunto de datos "MC Maze" (140 neuronas, 1720 ensayos), el CASSM logra puntuaciones de co-suavizado (co-smoothing) comparables a las de LFADS.
• Conjunto de Datos de Pez Cebra: En un conjunto de datos de gran escala con aproximadamente 88,000 neuronas y solo 96 ensayos, el bGPFA falla debido a errores de falta de memoria (out-of-memory). El CASSM procesa los datos con éxito y, al utilizar priors espaciales (CASSM-SP), logra un rendimiento estadísticamente similar al de LFADS mientras proporciona una incertidumbre calibrada.

Significancia y Reivindicaciones
El artículo afirma que el CASSM proporciona una "hoja de ruta" para que los investigadores elijan modelos de variables latentes dinámicas basados en las propiedades del conjunto de datos. Los autores argumentan que, mientras que los métodos de aprendizaje profundo (LFADS) son superiores cuando los ensayos superan a las neuronas, los métodos bayesianos se vuelven esenciales en el régimen de desequilibrio de escala típico de la neurociencia a gran escala. El CASSM cierra esta brecha ofreciendo un método que es computacionalmente escalable como las redes profundas, pero que retiene los sesgos inductivos y la cuantificación de la incertidumbre de la inferencia bayesiana. Los autores señalan modestamente que el CASSM está altamente especializado para este régimen específico y es un método lineal, lo que significa que podría tener dificultades si la dinámica subyacente es altamente no lineal o si el rango de la aproximación se elige incorrectamente. Sin embargo, en el régimen objetivo, ofrece una alternativa competitiva a las redes profundas ávidas de datos con una supervisión de la incertidumbre superior.