Computation-Aware Kalman Filtering with Model Selection for Neural Dynamics

Resumo Técnico: Filtragem de Kalman Consciente de Computação com Seleção de Modelo para Dinâmicas Neurais

Enunciado do Problema
O artigo aborda os desafios computacionais e estatísticos de modelar variáveis latentes dinâmicas em registros neurais modernos e de larga escala. Embora métodos Bayesianos (ex: Análise de Fator de Processo Gaussiano, GPFA) ofereçam priors explícitos e quantificação de incerteza, sua complexidade cúbica em relação à resolução temporal e complexidade quadrática em relação à resolução espacial (número de neurônios) os torna inviáveis para conjuntos de dados modernos. Por outro lado, modelos de aprendizado profundo superparametrizados (ex: LFADS) escalam favoravelmente, mas frequentemente carecem de uma quantificação de incerteza robusta e exigem um grande número de ensaios para superar sua falta de viés indutivo.

Os autores identificam um regime específico de "desequilíbrio de escala" prevalente na neurociência (ex: estudos do cérebro humano), onde o número de neurônios registrados excede significativamente o número de ensaios experimentais. Neste regime, os modelos padrão de aprendizado profundo degradam-se devido à escasseza de dados, enquanto as aproximações Bayesianas tradicionais frequentemente falham ao não considerar os erros introduzidos por suas próprias restrições computacionais, levando a estimativas de incerteza excessivamente confiantes ("privadas de variância").

Metodologia: Modelo de Espaço de Estados Consciente de Computação (CASSM)
Os autores propõem o Modelo de Espaço de Estados Consciente de Computação (CASSM), um framework numérico probabilístico que estende o Filtro de Kalman Consciente de Computação (CAKF) para incluir a seleção de modelo. A metodologia central envolve:

1. Projeção e Incerteza Computacional: O CASSM reconhece que o posterior "verdadeiro" é computacionalmente inviável de ser calculado diretamente. Em vez de ignorar o erro de aproximação, ele modela explicitamente a "incerteza computacional". O método projeta observações de alta dimensão ($y_k \in \mathbb{R}^n$) para um espaço de menor dimensão ($\check{y}_k \in \mathbb{R}^{\check{n}}$) usando uma matriz de política aprendida $S_k$. Isso reduz a complexidade da inversão da matriz de inovação de $O(n^3)$ para $O(\check{n}^3)$.
2. Atualizações de Covariância de Baixo Posto: Para lidar com restrições de memória ($O(d^2)$ para covariâncias de estado), o CASSM utiliza uma estrutura de downdate de baixo posto. Ele armazena raízes quadradas "esbeltas" de baixo posto das atualizações de covariância em vez das matrizes completas, reduzindo os requisitos de memória para $O(dr)$.
3. Seleção de Modelo via Análogo de PCA: A política de projeção $S_k$ é interpretada através de uma lente informacional. O Teorema 1 estabelece que a ação ótima (projeção) que minimiza a entropia do posterior latente corresponde aos autovetores superiores da covariância de inovação (análogo à Análise de Componentes Principais - PCA). O framework aprende essas ações para alinhar-se com as direções de maior incerteza esperada.
4. Objetivo de Treinamento Inovador: Os autores introduzem uma função de perda específica derivada do Limite Inferior de Evidência (ELBO). Diferente da Verossimilhança Negativa de Logaritmo (NLL) padrão, que pode degradar sob projeção, o objetivo do CASSM equilibra um termo de ajuste de dados com um termo de divergência. Este termo de divergência penaliza o desvio do posterior aproximado em relação à distribuição de filtragem exata, garantindo que as estimativas de incerteza se ampliem apropriadamente conforme a aproximação computacional aumenta.
5. Priors Espaciais: Ao traduzir representações de Processos Gaussianos para Modelos de Espaço de Estados, o CASSM permite a codificação explícita de priors espaciais (ex: coordenadas de neurônios ou tipos celulares) via um kernel espacial $\Sigma^{(x)}$, uma característica que é apenas implícita em métodos anteriores como GPFA ou LFADS.
6. Pré-processamento e Estabilidade: O pipeline inclui transformações estabilizadoras de variância (Anscombe) para dados de disparos (spikes) e utiliza decomposição de valor singular diferenciável (baseada na pseudoinversa de Moore-Penrose) para garantir a estabilidade numérica durante a truncagem das atualizações de baixo posto.

Principais Contribuições

• Framework CASSM: Um método numérico probabilístico escalável projetado especificamente para o regime de desequilíbrio de escala (ensaios $\ll$ neurônios).
• Inferência Tratável com Calibração de Incerteza: O método produz inferência tratável em espaços de estados grandes, fornecendo simultaneamente uma incerteza calibrada que considera os erros de aproximação computacional, evitando a excessiva confiança típica de outras aproximações.
• Função de Perda Inovadora: Uma função de perda de treinamento que permite a otimização eficiente de hiperparâmetros e seleção de modelo ao penalizar a divergência da distribuição de filtragem exata, permitindo que o modelo aprenda projeções de dados informativas.
• Modelagem Espacial Explícita: Uma técnica para incorporar explicitamente priors espaciais sobre neurônios dentro de um modelo de espaço de estados, diferenciando-o de modelos de variáveis latentes que tratam a estrutura espacial de forma implícita.

Resultos
Os autores avaliam o CASSM tanto em dados sintéticos quanto em dados do mundo real:

• Dados de Lorenz Sintéticos: O CASSM demonstra escalonamento linear em relação ao número de neurônios, superando o GPFA e o bGPFA em tempo computacional em grandes escalas ($10^4$ neurônios). No regime de desequilíbrio de escala, o CASSM mantém um desempenho preditivo competitivo (MSE) comparado ao LFADS, enquanto o LFADS degrada rapidamente quando os ensaios são limitados.
• Calibração de Incerteza: Experimentos de cobertura mostram que o CASSM fornece intervalos de confiança de 95% significativamente melhor calibrados em comparação ao bGPFA e GPFA, que tendem a ser excessivamente confiantes.
• Dados Neurais Reais (Neural Latents Benchmark): Em conjuntos de dados como Area 2 Bump e DMFC, o CASSM performa competitivamente com o LFADS e bGPFA. No conjunto de dados "MC Maze" (140 neurônios, 1720 ensaios), o CASSM alcança pontuações de co-suavização (co-smoothing) comparáveis às do LFADS.
• Conjunto de Dados de Zebrafish: Em um conjunto de dados de larga escala com ~88.000 neurônios e apenas 96 ensaios, o bGPFA falha devido a erros de falta de memória (out-of-memory). O CASSM processa os dados com sucesso e, ao utilizar priors espaciais (CASSM-SP), alcança um desempenho estatisticamente similar ao LFADS, fornecendo incerteza calibrada.

Significância e Alegações
O artigo afirma que o CASSM fornece um "roteiro" para pesquisadores escolherem modelos de variáveis latentes dinâmicas baseados nas propriedades do conjunto de dados. Os autores argumentam que, enquanto os métodos de aprendizado profundo (LFADS) são superiores quando os ensaios superam os neurônios, os métodos Bayesianos tornam-se essenciais no regime de desequilíbrio de escala, típico da neurociência de larga escala. O CASSM preenche essa lacção ao oferecer um método que é computacionalmente escalável como redes profundas, mas que retém os vieses indutivos e a quantificação de incerteza da inferência Bayesiana. Os autores observam modestamente que o CASSM é altamente especializado para este regime específico e é um método linear, o que significa que pode ter dificuldades se a dinâmica subjacente for altamente não-linear ou se o posto da aproximação for mal escolhido. No entanto, no regime alvo, ele oferece uma alternativa competitiva a redes profundas ávidas por dados, com uma calibração de incerteza superior.