Resumo Técnico: CaMBRAIN - Inferência Contínua em Tempo Real de EEG com Modelos de Espaço de Estados Causais

Declaração do Problema

Os sinais de Eletroencefalografia (EEG) são críticos para monitorar a atividade cerebral em ambientes clínicos e biomecânicos, mas apresentam desafios únicos para a aprendizagem profunda devido ao seu comprimento ilimitado, frequentemente abrangendo minutos a horas. As abordagens existentes de aprendizagem profunda, particularmente aquelas baseadas em Transformers, enfrentam dois obstáculos principais ao serem aplicadas a EEG de longa duração:

1. Ineficiência Computacional: Mecanismos de atenção incorrem em escalamento quadrático ($O(T^2)$) com o comprimento da sequência, tornando-os ineficientes para sinais longos.
2. Limitações de Contexto: Para gerenciar custos computacionais, os modelos atuais geralmente processam EEG por meio de janelas deslizantes com comprimentos de entrada fixos. Essa abordagem força o modelo a recalcular características sobre regiões sobrepostas e impede a manutenção de uma representação persistente de todo o histórico do sinal. Consequentemente, esses modelos carecem de compreensão temporal global e não podem realizar inferência contínua em tempo real verdadeira, onde cada nova observação se baseia no contexto anterior sem reexaminar entradas passadas.

Além disso, os objetivos existentes de pré-treinamento auto-supervisionado para EEG frequentemente dependem de reconstrução de sinal (por exemplo, autoencoders mascarados). Embora eficazes para aprender a fidelidade local do sinal, esses objetivos não treinam explicitamente o estado oculto para reter contexto de longo alcance saliente, o que é essencial para inferência em fluxo contínuo, onde o modelo deve comprimir o histórico em uma memória compacta.

Metodologia

Os autores propõem o CaMBRAIN, o primeiro modelo de EEG baseado em Mamba e causal, projetado para inferência contínua em fluxo em tempo real. A metodologia é construída sobre dois pilares centrais: uma arquitetura causal e um pipeline de treinamento auto-supervisionado especializado em múltiplos estágios.

1. Arquitetura: Mamba-3 Causal

O CaMBRAIN utiliza um Modelo de Espaço de Estados (SSM) unidirecional e causal, especificamente a arquitetura Mamba-3, para processar sinais de EEG.

• Processamento de Entrada: Sinais de EEG brutos são particionados em patches não sobrepostos de 62,5 ms (16 amostras a 256 Hz) em todos os canais. Esses patches são incorporados em tokens usando um incorporador de canais baseado em transformer.
• Estado Oculto Persistente: Diferentemente das abordagens de janela deslizante, o CaMBRAIN mantém um estado oculto persistente que comprime incrementalmente todo o contexto passado. À medida que novos patches chegam, o estado oculto é atualizado recursivamente ($h_t = \bar{A}h_{t-1} + \bar{B}x_t$) sem reexaminar entradas anteriores.
• Restrição Causal: O modelo opera estritamente de maneira apenas para frente, garantindo que as previsões no tempo $t$ dependam apenas de entradas passadas e atuais ($x_{\le t}$), alinhando-se com a realidade física do monitoramento de EEG.
• Transições de Valor Complexo: Aproveitando o Mamba-3, o modelo emprega transições de estado de valor complexo para capturar melhor as dinâmicas oscilatórias e as relações de fase inerentes a sinais de EEG de múltiplas bandas.

2. Pipeline de Treinamento

Para abordar as limitações de objetivos baseados em reconstrução para fluxo contínuo, os autores introduzem um pipeline de treinamento de três estágios:

• Estágio 1: Pré-treinamento Preditivo Causal:
  O codificador é pré-treinado usando uma perda causal híbrida para instilar estrutura temporal geral. Isso combina:

  • Modelagem Autoregressiva (ARM): Prever a incorporação do próximo patch a partir do estado oculto atual.
  • Reconstrução Mascarada: Reconstruir tokens de EEG brutos mascarados a partir do contexto causal fornecido pelo codificador.
    Este estágio garante que o modelo aprenda dinâmicas temporais locais e fidelidade de sinal.

• Estágio 2: Refinamento de Representação Latente sem Reconstrução (Estilo JEPA):
  Para incentivar a retenção de informações preditivas de longo alcance, o modelo passa por uma fase de treinamento sem decodificador inspirada em Arquiteturas de Incorporação Conjunta Preditiva (JEPA).

  • Framework Estudante-Professor: Uma rede estudante processa uma visão mascarada da entrada, enquanto uma rede professora (uma Média Móvel Exponencial da estudante) processa a entrada não mascarada.
  • Objetivos: A estudante é treinada para prever:
    1. Latenas Mascaradas: As representações da professora para posições mascaradas.
    2. Latenas Futuras: Representações futuras de múltiplos passos ($K$ passos à frente) a partir do estado oculto atual.
       Este estágio força explicitamente o estado oculto a codificar características informativas, invariantes e temporalmente relevantes necessárias para inferência contínua, em vez de apenas reconstruir detalhes locais do sinal.

• Estágio 3: Ajuste Fino para Tarefas Específicas:
  O modelo pré-treinado é ajustado finamente em tarefas específicas (por exemplo, detecção de convulsões, classificação de artefatos) usando aprendizagem supervisionada padrão. Notavelmente, o CaMBRAIN utiliza apenas o estado oculto final para classificação, evitando a necessidade de agrupamento ou agregação de tokens através de segmentos fixos.

Contribuições Principais

1. Formulação de Inferência Causal: O artigo propõe uma formulação de inferência de EEG unidirecional que se afasta das abordagens de janela deslizante e transformers bidirecionais, tratando sinais como fluxos contínuos em vez de segmentos independentes.
2. Arquitetura CaMBRAIN: Um codificador Mamba-3 causal que substitui o recálculo redundante de janela deslizante por um estado oculto persistente, permitindo inferência eficiente e escalável sobre sinais arbitrariamente longos com complexidade de tempo linear.
3. Treinamento em Nível de Representação: Um novo framework auto-supervisionado de múltiplos estágios que vai além da reconstrução de sinal. Ao combinar previsão autoregressiva com previsão de latentes mascarados e futuros, ele treina explicitamente o estado oculto para reter contexto de longo alcance saliente.

Resultados

O CaMBRAIN foi avaliado em quatro conjuntos de dados de EEG diversos: TUAR (detecção de artefatos), TUAB (detecção de EEG anormal), MAT (detecção de estresse mental) e CHB-MIT (detecção de convulsões).

• Desempenho: O CaMBRAIN alcançou resultados de última geração (SOTA) ou competitivos em todos os benchmarks.
  • TUAR: Alcançou uma AUROC de 0,936, superando modelos de base maiores como o LUNA-H (0,921).
  • CHB-MIT: Alcançou uma AUROC de 0,921, demonstrando forte desempenho em dados de convulsões altamente desbalanceados.
  • MAT: Alcançou a melhor acurácia balanceada (0,778) e AUROC (0,876) entre os métodos comparados.
• Eficiência Computacional: O modelo demonstrou ganhos significativos de eficiência em cenários de fluxo contínuo.
  • Em uma taxa de atualização de fluxo contínuo de 16 Hz, o CaMBRAIN exigiu 1,23 GFLOPs/s.
  • Isso representa uma redução de 6,1× a 158× no cálculo sustentado em comparação com outras bases de modelos de fundação (por exemplo, o REVE exigiu 194,79 GFLOPs/s), enquanto igualava ou superava sua precisão.
• Estudos de Ablação:
  • Estado Persistente: Usar um estado oculto persistente (vs. redefinir a cada 5 segundos) melhorou a AUROC em +0,015 no CHB-MIT e melhorou significativamente a localização do início da convulsão.
  • Estágios de Treinamento: O pré-treinamento completo de dois estágios (Estágio 1 + Estágio 2) produziu ganhos substancialmente maiores (+2,30 AUROC) sobre o ajuste fino padrão ou o Estágio 1 sozinho, confirmando a necessidade de refinamento em nível de representação para tarefas de fluxo contínuo.

Significado e Alegações

O artigo alega que o CaMBRAIN representa uma mudança de paradigma na modelagem de EEG, alinhando tanto a arquitetura quanto os objetivos de treinamento com a natureza causal e contínua inerente aos sinais de EEG.

• Capacidade em Tempo Real: Permite o primeiro modelo capaz de realizar inferência de longo alcance e contínua em sinais de EEG de comprimento variável, sem a penalidade computacional de escalamento quadrático ou a redundância de janelas deslizantes.
• Relevância Clínica: A abordagem espelha como os clínicos monitoram EEG à beira do leito, mantendo uma memória compacta de todo o histórico do sinal para detectar eventos que podem estar separados por longos intervalos.
• Eficiência: Ao aproveitar o escalamento de tempo linear dos Modelos de Espaço de Estados e um objetivo de treinamento em nível de representação, o CaMBRAIN alcança alta vazão e baixa latência, tornando-o adequado para aplicações clínicas e vestíveis em tempo real onde os recursos computacionais são limitados.

Os autores reconhecem que, embora o CaMBRAIN estabeleça novos benchmarks em vários conjuntos de dados, ele permanece competitivo, e não superior, na detecção de EEG anormal do TUAB em comparação com os modelos de fundação mais fortes, sugerindo que a escalabilidade do corpus de pré-treinamento para corresponder a conjuntos de dados maiores e mais heterogêneos (como o REVE) é uma direção futura necessária.