Resumen Técnico: CaMBRAIN - Inferencia Continua de EEG en Tiempo Real con Modelos de Espacio de Estados Causales

Enunciado del Problema

Las señales de electroencefalografía (EEG) son críticas para monitorear la actividad cerebral en entornos clínicos y biomecánicos, pero presentan desafíos únicos para el aprendizaje profundo debido a su longitud ilimitada, que a menudo abarca desde minutos hasta horas. Los enfoques existentes de aprendizaje profundo, particularmente aquellos basados en Transformers, enfrentan dos obstáculos principales al aplicarse a EEG de larga duración:

1. Ineficiencia Computacional: Los mecanismos de atención incurrir en una escalabilidad cuadrática ($O(T^2)$) con la longitud de la secuencia, lo que los hace ineficientes para señales largas.
2. Limitaciones de Contexto: Para gestionar los costos computacionales, los modelos actuales procesan el EEG mediante ventanas deslizantes con longitudes de entrada fijas. Este enfoque obliga al modelo a recalcular características sobre regiones superpuestas e impide el mantenimiento de una representación persistente de todo el historial de la señal. En consecuencia, estos modelos carecen de una comprensión temporal global y no pueden realizar una inferencia continua y en tiempo real donde cada nueva observación se construya sobre el contexto previo sin reexaminar las entradas pasadas.

Además, los objetivos existentes de preentrenamiento auto-supervisado para EEG a menudo dependen de la reconstrucción de señales (por ejemplo, autoencoders enmascarados). Si bien son efectivos para aprender la fidelidad local de la señal, estos objetivos no entrenan explícitamente el estado oculto para retener el contexto de largo alcance saliente, lo cual es esencial para la inferencia en flujo continuo donde el modelo debe comprimir el historial en una memoria compacta.

Metodología

Los autores proponen CaMBRAIN, el primer modelo de EEG basado en Mamba y causal diseñado para inferencia en flujo continuo en tiempo real. La metodología se basa en dos pilares fundamentales: una arquitectura causal y una tubería de entrenamiento auto-supervisado especializada de múltiples etapas.

1. Arquitectura: Mamba-3 Causal

CaMBRAIN utiliza un Modelo de Espacio de Estados (SSM) unidireccional y causal, específicamente la arquitectura Mamba-3, para procesar señales de EEG.

• Procesamiento de Entrada: Las señales de EEG crudas se dividen en parches no superpuestos de 62.5 ms (16 muestras a 256 Hz) a través de todos los canales. Estos parches se incrustan en tokens utilizando un incrustador de canales basado en transformers.
• Estado Oculto Persistente: A diferencia de los enfoques de ventana deslizante, CaMBRAIN mantiene un estado oculto persistente que comprime incrementalmente todo el contexto pasado. A medida que llegan nuevos parches, el estado oculto se actualiza recurrentemente ($h_t = \bar{A}h_{t-1} + \bar{B}x_t$) sin reexaminar las entradas anteriores.
• Restricción Causal: El modelo opera estrictamente de manera solo hacia adelante, asegurando que las predicciones en el tiempo $t$ dependan únicamente de las entradas pasadas y actuales ($x_{\le t}$), alineándose con la realidad física del monitoreo de EEG.
• Transiciones de Valor Complejo: Aprovechando Mamba-3, el modelo emplea transiciones de estado de valor complejo para capturar mejor las dinámicas oscilatorias y las relaciones de fase inherentes a las señales de EEG de múltiples bandas.

2. Tubería de Entrenamiento

Para abordar las limitaciones de los objetivos basados en reconstrucción para el flujo continuo, los autores introducen una tubería de entrenamiento de tres etapas:

• Etapa 1: Preentrenamiento Predictivo Causal:
  El codificador se preentrena utilizando una pérdida causal híbrida para instilar una estructura temporal general. Esto combina:

  • Modelado Autoregresivo (ARM): Predecir la siguiente incrustación de parche a partir del estado oculto actual.
  • Reconstrucción Enmascarada: Reconstruir tokens de EEG crudos enmascarados a partir del contexto causal proporcionado por el codificador.
    Esta etapa asegura que el modelo aprenda dinámicas temporales locales y fidelidad de la señal.

• Etapa 2: Refinamiento de Representación Latente sin Reconstrucción (Estilo JEPA):
  Para fomentar la retención de información predictiva de largo alcance, el modelo se somete a una fase de entrenamiento sin decodificador inspirada en Arquitecturas Predictivas de Incrustación Conjunta (JEPA).

  • Marco Estudiante-Profesor: Una red estudiante procesa una vista enmascarada de la entrada, mientras que una red profesora (un Promedio Móvil Exponencial de la estudiante) procesa la entrada sin enmascarar.
  • Objetivos: La estudiante se entrena para predecir:
    1. Latenas Enmascaradas: Las representaciones del profesor para posiciones enmascaradas.
    2. Latenas Futuras: Representaciones futuras de múltiples pasos ($K$ pasos adelante) desde el estado oculto actual.
       Esta etapa fuerza explícitamente al estado oculto a codificar características informativas, invariantes y temporalmente relevantes necesarias para la inferencia continua, en lugar de simplemente reconstruir detalles locales de la señal.

• Etapa 3: Ajuste Fino para Tareas Específicas:
  El modelo preentrenado se ajusta fino en tareas específicas (por ejemplo, detección de convulsiones, clasificación de artefactos) utilizando aprendizaje supervisado estándar. Cabe destacar que CaMBRAIN utiliza únicamente el estado oculto final para la clasificación, evitando la necesidad de agrupar o agregar tokens a través de segmentos fijos.

Contribuciones Clave

1. Formulación de Inferencia Causal: El artículo propone una formulación de inferencia de EEG unidireccional que se aparta de los enfoques de ventana deslizante y transformers bidireccionales, tratando las señales como flujos continuos en lugar de segmentos independientes.
2. Arquitectura CaMBRAIN: Un codificador Mamba-3 causal que reemplaza el recálculo redundante de ventanas deslizantes con un estado oculto persistente, habilitando una inferencia eficiente y escalable sobre señales arbitrariamente largas con complejidad de tiempo lineal.
3. Entrenamiento a Nivel de Representación: Un marco auto-supervisado novedoso de múltiples etapas que va más allá de la reconstrucción de señales. Al combinar la predicción autoregresiva con la predicción de latentes enmascarados y futuros, entrena explícitamente al estado oculto para retener el contexto de largo alcance saliente.

Resultados

CaMBRAIN fue evaluado en cuatro conjuntos de datos de EEG diversos: TUAR (detección de artefactos), TUAB (detección de EEG anormal), MAT (detección de estrés mental) y CHB-MIT (detección de convulsiones).

• Rendimiento: CaMBRAIN logró resultados de vanguardia (SOTA) o competitivos en todos los puntos de referencia.
  • TUAR: Logró un AUROC de 0.936, superando a modelos fundacionales más grandes como LUNA-H (0.921).
  • CHB-MIT: Logró un AUROC de 0.921, demostrando un rendimiento sólido en datos de convulsiones altamente desequilibrados.
  • MAT: Logró la mejor precisión equilibrada (0.778) y AUROC (0.876) entre los métodos comparados.
• Eficiencia Computacional: El modelo demostró ganancias significativas de eficiencia en escenarios de flujo continuo.
  • A una tasa de actualización de flujo continuo de 16 Hz, CaMBRAIN requirió 1.23 GFLOPs/s.
  • Esto representa una reducción de 6.1× a 158× en cómputo sostenido en comparación con otras líneas base de modelos fundacionales (por ejemplo, REVE requirió 194.79 GFLOPs/s), mientras iguala o supera su precisión.
• Estudios de Ablación:
  • Estado Persistente: El uso de un estado oculto persistente (frente a reiniciar cada 5 segundos) mejoró el AUROC en +0.015 en CHB-MIT y mejoró significativamente la localización del inicio de convulsiones.
  • Etapas de Entrenamiento: El preentrenamiento completo de dos etapas (Etapa 1 + Etapa 2) produjo ganancias sustancialmente mayores (+2.30 AUROC) sobre el ajuste fino estándar o la Etapa 1 sola, confirmando la necesidad de refinamiento a nivel de representación para tareas de flujo continuo.

Significado y Afirmaciones

El artículo afirma que CaMBRAIN representa un cambio de paradigma en el modelado de EEG al alinear tanto la arquitectura como los objetivos de entrenamiento con la naturaleza causal y continua inherente de las señales de EEG.

• Capacidad en Tiempo Real: Habilita el primer modelo capaz de realizar inferencia continua de largo alcance en señales de EEG de longitud variable sin la penalización computacional de la escalabilidad cuadrática o la redundancia de las ventanas deslizantes.
• Relevancia Clínica: El enfoque refleja cómo los clínicos monitorean el EEG al pie de la cama, manteniendo una memoria compacta de todo el historial de la señal para detectar eventos que pueden estar separados por largos intervalos.
• Eficiencia: Al aprovechar la escalabilidad de tiempo lineal de los Modelos de Espacio de Estados y un objetivo de entrenamiento a nivel de representación, CaMBRAIN logra un alto rendimiento y baja latencia, haciéndolo adecuado para aplicaciones clínicas y portátiles en tiempo real donde los recursos computacionales son limitados.

Los autores reconocen que, aunque CaMBRAIN establece nuevos puntos de referencia en varios conjuntos de datos, sigue siendo competitivo en lugar de superior en la detección de EEG anormal en TUAB en comparación con los modelos fundacionales más fuertes, lo que sugiere que escalar el corpus de preentrenamiento para igualar conjuntos de datos más grandes y heterogéneos (como REVE) es una dirección futura necesaria.