OmniMouse: Scaling properties of multi-modal, multi-task Brain Models on 150B Neural Tokens

1. Problema e Motivação

O avanço recente na Inteligência Artificial, impulsionado pela escalabilidade de dados e modelos (como em LLMs e modelos de visão computacional), levantou a questão de se princípios semelhantes se aplicariam à modelagem da atividade cerebral. Até o momento, os modelos de neurociência enfrentam desafios significativos:

• Dados Fragmentados: Os conjuntos de dados neurais são geralmente menores, mais fragmentados e menos diversos do que os corpora de internet usados em IA.
• Arquiteturas Especializadas: Modelos anteriores focavam em modalidades isoladas (apenas neurônios, apenas fMRI, apenas EEG) ou tarefas únicas, sem a capacidade de unificar previsão neural, decodificação comportamental e previsão de estímulos em uma única arquitetura.
• Incerteza sobre Leis de Escala: Não havia consenso sobre se as "leis de escala" (onde o aumento de dados e parâmetros melhora o desempenho) se aplicariam a dados neurais ruidosos e complexos.

O objetivo deste trabalho é introduzir uma arquitetura unificada capaz de modelar a atividade do córtex visual de camundongos em escala massiva e investigar como o aumento de dados e tamanho do modelo afeta o desempenho.

2. Metodologia

A. O Conjunto de Dados (OmniMouse Dataset)

Os autores compilaram um dos maiores conjuntos de dados de neurônios únicos já utilizados:

• Escala: Mais de 3,1 milhões de neurônios de 73 camundongos (em 323 sessões).
• Tokens: Totalizando mais de 150 bilhões de tokens neurais.
• Modalidades:
  • Neuronal: Gravações de imagem de cálcio de dois fótons (6-14 Hz) no córtex visual (V1 e áreas visuais superiores) de camundongos acordados e comportando-se.
  • Estímulos: Vídeos naturalísticos (filmes, Sports-1M), imagens (ImageNet) e estímulos paramétricos (Gabor, ruído, etc.).
  • Comportamento: Velocidade de corrida, dilatação pupilar, posição da pupila e sua derivada.
• Pré-processamento: Os dados foram alinhados temporalmente, com comportamentos e estímulos reamostrados para 20-30 Hz e respostas neurais interpoladas para 30 Hz.

B. Arquitetura OmniMouse

O OmniMouse é um modelo multimodal e multitarefa baseado em Transformers, projetado para flexibilidade extrema através de mascaramento estruturado.

1. Tokenização:

   • Respostas Neurais: Cada neurônio é tokenizado usando convolução 1D com stride, enriquecido com embeddings de identidade (neurônio, sessão, animal).
   • Estímulos Visuais: Utiliza um encoder Hiera (Vision Transformer hierárquico) para extrair embeddings espaço-temporais de vídeos.
   • Comportamento: Traços comportamentais são projetados linearmente com embeddings de canal e identidade.

2. Mecanismo de Mascaramento (Flexibilidade):

   • O modelo aceita combinações arbitrárias de entradas e saídas. O mascaramento define quais tokens são codificados (contexto) e quais são reconstruídos (alvo).
   • Tarefas Suportadas:
     • Previsão Neural (Forecasting): Prever atividade futura baseada no passado.
     • Decodificação Comportamental: Prever comportamento baseado na atividade neural.
     • Codificação de Estímulo: Prever resposta neural baseada no vídeo.
     • Previsão de Subpopulação: Prever neurônios não observados baseando-se em uma subpopulação observada.
     • Combinações: Qualquer mistura das acima (ex: prever neurônios dados vídeo + comportamento + histórico neural).

3. Arquitetura do Modelo:

   • Encoder: Comprime tokens não mascarados em um latente fixo usando atenção cruzada com queries aprendidas.
   • Fusão Multimodal: Uma pilha de camadas Transformer que integra latentes neurais/comportamentais com características visuais. Utiliza atenção deslizante local (para eficiência) e camadas globais (para contexto de longo alcance).
   • Decoder: Reconstrói os tokens alvo (neurônios ou comportamento) usando atenção cruzada causal sobre a representação fundida.

C. Treinamento e Escala

• Função de Perda: Perda de Poisson para respostas neurais e Erro Quadrático Médio (MSE) para comportamento (com peso reduzido de 0.1 para o comportamento).
• Estratégia de Escala: Os autores treinaram modelos variando de 1M a 300M parâmetros (escalonando largura e profundidade) e variando o tamanho do conjunto de dados de 8 a 323 sessões.
• Protocolo de Avaliação: Uso de 7 camundongos de teste (incluindo dados públicos do SENSORIUM 2022/2023) para comparação justa com baselines especializados.

3. Principais Contribuições e Resultados

A. Desempenho State-of-the-Art (SOTA)

O OmniMouse superou consistentemente todas as baselines especializadas (como MtM, CEBRA, POYO+, Neuroformer) em quase todas as tarefas, tanto em cenários de dados combinados quanto em comparações com dados idênticos (data-matched).

• SENSORIUM: Estabeleceu novos recordes nas competições SENSORIUM 2022 e 2023, superando os vencedores anteriores mesmo com backbones congelados e apenas ajuste fino de embeddings.

B. Descoberta sobre Leis de Escala (Scaling Laws)

O estudo revelou uma inversão interessante em relação à IA tradicional (LLMs):

1. Dados são o Gargalo: O desempenho melhora sistematicamente e de forma confiável com o aumento da quantidade de dados (sessões de gravação).
2. Saturação de Parâmetros: Ao contrário dos modelos de linguagem, onde aumentar o tamanho do modelo continua gerando ganhos, o desempenho do OmniMouse satura em tamanhos moderados (acima de ~80M parâmetros).
   • Implicação: Modelos de cérebro atuais são limitados pela disponibilidade de dados, não pela capacidade computacional ou número de parâmetros. Mesmo em um sistema "simples" como o córtex visual de camundongo, os modelos ainda não exploraram todo o potencial dos dados disponíveis.

C. Dinâmica de Escala por Tarefa

• Decodificação Comportamental: Mostrou os maiores ganhos com o aumento do tamanho do modelo, sugerindo que tarefas comportamentais podem ainda se beneficiar de modelos maiores (não saturou completamente).
• Previsão Neural: Os maiores ganhos vieram do aumento dos dados, especialmente em tarefas que envolvem entrada de vídeo.

4. Significado e Impacto

• Fundação para Neurociência: O OmniMouse demonstra que é possível construir modelos de base (foundation models) para neurociência que unificam múltiplas modalidades e tarefas, superando abordagens fragmentadas.
• Mudança de Paradigma na Escala: A descoberta de que os ganhos saturam com o tamanho do modelo, mas não com os dados, sugere que o futuro do progresso em modelagem cerebral depende da coleta de dados em larga escala e diversificada, e não apenas de aumentar a arquitetura.
• Potencial para Transições de Fase: A observação de escalabilidade sistemática levanta a possibilidade de que, com conjuntos de dados ainda maiores e mais ricos (talvez de múltiplas áreas cerebrais ou espécies), os modelos possam atingir "transições de fase" e adquirir capacidades qualitativamente novas, semelhantes às emergentes em LLMs.
• Reprodutibilidade: O código e os dados (em preparação para liberação pública) são disponibilizados, permitindo que a comunidade valide e expanda este trabalho.

Conclusão

O OmniMouse representa um marco na modelagem de dados neurais, provando que a unificação de tarefas e modalidades em uma arquitetura escalável é viável e superior. O trabalho redefine a estratégia de pesquisa em neurociência computacional, indicando que a prioridade deve ser a expansão de conjuntos de dados multimodais massivos para desbloquear o verdadeiro potencial dos modelos de cérebro, em vez de focar exclusivamente no aumento de parâmetros.