Riepilogo Tecnico: CaMBRAIN - Inferenza EEG Continua e in Tempo Reale con Modelli a Spazio di Stato Causale

Enunciato del Problema

I segnali di elettroencefalografia (EEG) sono fondamentali per il monitoraggio dell'attività cerebrale in contesti clinici e biomeccanici, ma presentano sfide uniche per l'apprendimento profondo a causa della loro lunghezza illimitata, che spesso si estende da minuti a ore. Gli approcci esistenti basati sull'apprendimento profondo, in particolare quelli fondati sui Transformer, affrontano due ostacoli principali quando applicati all'EEG a lunga durata:

1. Inefficienza Computazionale: I meccanismi di attenzione comportano una scalatura quadratica ($O(T^2)$) rispetto alla lunghezza della sequenza, rendendoli inefficienti per segnali lunghi.
2. Limitazioni del Contesto: Per gestire i costi computazionali, i modelli attuali elaborano tipicamente l'EEG tramite finestre scorrevoli con lunghezze di input fisse. Questo approccio costringe il modello a ricalcolare le caratteristiche su regioni sovrapposte e impedisce il mantenimento di una rappresentazione persistente dell'intera storia del segnale. Di conseguenza, questi modelli mancano di una comprensione temporale globale e non possono eseguire una vera inferenza continua e in tempo reale in cui ogni nuova osservazione si basa sul contesto precedente senza riesaminare gli input passati.

Inoltre, gli obiettivi esistenti di pre-addestramento auto-supervisionato per l'EEG spesso si basano sulla ricostruzione del segnale (ad esempio, autoencoder mascherati). Sebbene efficaci per apprendere la fedeltà locale del segnale, questi obiettivi non addestrano esplicitamente lo stato nascosto a mantenere un contesto a lungo raggio saliente, essenziale per l'inferenza in streaming, dove il modello deve comprimere la storia in una memoria compatta.

Metodologia

Gli autori propongono CaMBRAIN, il primo modello EEG basato su Mamba e causale, progettato per l'inferenza continua in streaming in tempo reale. La metodologia si fonda su due pilastri principali: un'architettura causale e una pipeline di addestramento auto-supervisionato multistadio specializzata.

1. Architettura: Causal Mamba-3

CaMBRAIN utilizza un Modello a Spazio di Stato (SSM) unidirezionale e causale, specificamente l'architettura Mamba-3, per elaborare i segnali EEG.

• Elaborazione dell'Input: I segnali EEG grezzi sono suddivisi in patch non sovrapposte di 62,5 ms (16 campioni a 256 Hz) su tutti i canali. Queste patch sono incorporate in token utilizzando un incorporatore di canali basato su transformer.
• Stato Nascosto Persistente: A differenza degli approcci a finestra scorrevole, CaMBRAIN mantiene uno stato nascosto persistente che comprime in modo incrementale tutto il contesto passato. Man mano che arrivano nuove patch, lo stato nascosto viene aggiornato in modo ricorsivo ($h_t = \bar{A}h_{t-1} + \bar{B}x_t$) senza riesaminare gli input precedenti.
• Vincolo Causale: Il modello opera strettamente in modalità solo in avanti, garantendo che le previsioni al tempo $t$ dipendano solo dagli input passati e correnti ($x_{\le t}$), in linea con la realtà fisica del monitoraggio EEG.
• Transizioni a Valore Complesso: Sfruttando Mamba-3, il modello impiega transizioni di stato a valore complesso per catturare meglio le dinamiche oscillatorie e le relazioni di fase intrinseche nei segnali EEG multibanda.

2. Pipeline di Addestramento

Per affrontare le limitazioni degli obiettivi basati sulla ricostruzione per lo streaming, gli autori introducono una pipeline di addestramento in tre fasi:

• Fase 1: Pre-addestramento Predittivo Causale:
  Il codificatore viene pre-addestrato utilizzando una perdita causale ibrida per instillare una struttura temporale generale. Questa combina:

  • Modellazione Autoregressiva (ARM): Prevedere l'incorporazione della patch successiva dallo stato nascosto corrente.
  • Ricostruzione Mascherata: Ricostruire i token EEG grezzi mascherati dal contesto causale fornito dal codificatore.
    Questa fase garantisce che il modello apprenda le dinamiche temporali locali e la fedeltà del segnale.

• Fase 2: Rifinitura della Rappresentazione Latente Senza Ricostruzione (stile JEPA):
  Per incoraggiare la ritenzione di informazioni predittive a lungo raggio, il modello subisce una fase di addestramento senza decodificatore ispirata alle Architetture Predittive di Incorporamento Congiunto (JEPA).

  • Framework Studente-Insegnante: Una rete studente elabora una vista mascherata dell'input, mentre una rete insegnante (una Media Mobile Esponenziale dello studente) elabora l'input non mascherato.
  • Obiettivi: Lo studente viene addestrato a prevedere:
    1. Latenti Mascherati: Le rappresentazioni dell'insegnante per le posizioni mascherate.
    2. Latenti Futuri: Rappresentazioni future a più passi ($K$ passi avanti) dallo stato nascosto corrente.
       Questa fase forza esplicitamente lo stato nascosto a codificare caratteristiche informative, invarianti e temporalmente rilevanti necessarie per l'inferenza continua, piuttosto che limitarsi a ricostruire dettagli locali del segnale.

• Fase 3: Fine-tuning per Attività Specifiche:
  Il modello pre-addestrato viene affinato su compiti specifici (ad esempio, rilevamento di crisi epilettiche, classificazione di artefatti) utilizzando l'apprendimento supervisionato standard. Nota bene: CaMBRAIN utilizza solo lo stato nascosto finale per la classificazione, evitando la necessità di pooling o aggregazione di token su segmenti fissi.

Contributi Chiave

1. Formulazione dell'Inferenza Causale: Il documento propone una formulazione di inferenza EEG unidirezionale che si discosta dagli approcci a finestra scorrevole e bidirezionali basati su transformer, trattando i segnali come flussi continui piuttosto che segmenti indipendenti.
2. Architettura CaMBRAIN: Un codificatore Mamba-3 causale che sostituisce il ridondante ricalcolo a finestra scorrevole con uno stato nascosto persistente, abilitando un'inferenza efficiente e scalabile su segnali arbitrariamente lunghi con complessità temporale lineare.
3. Addestramento a Livello di Rappresentazione: Un nuovo framework auto-supervisionato multistadio che va oltre la ricostruzione del segnale. Combinando la previsione autoregressiva con la previsione di latenti mascherati e futuri, addestra esplicitamente lo stato nascosto a mantenere un contesto a lungo raggio saliente.

Risultati

CaMBRAIN è stato valutato su quattro dataset EEG diversificati: TUAR (rilevamento di artefatti), TUAB (rilevamento di EEG anomalo), MAT (rilevamento di stress mentale) e CHB-MIT (rilevamento di crisi epilettiche).

• Prestazioni: CaMBRAIN ha ottenuto risultati all'avanguardia (SOTA) o competitivi su tutti i benchmark.
  • TUAR: Ha raggiunto un AUROC di 0,936, superando modelli fondazionali più grandi come LUNA-H (0,921).
  • CHB-MIT: Ha raggiunto un AUROC di 0,921, dimostrando prestazioni solide su dati di crisi altamente sbilanciati.
  • MAT: Ha raggiunto la migliore accuratezza bilanciata (0,778) e AUROC (0,876) tra i metodi confrontati.
• Efficienza Computazionale: Il modello ha dimostrato guadagni significativi di efficienza negli scenari di streaming.
  • A un tasso di aggiornamento in streaming di 16 Hz, CaMBRAIN ha richiesto 1,23 GFLOPs/s.
  • Questo rappresenta una riduzione da 6,1× a 158× nel calcolo sostenuto rispetto ad altre basi di modelli fondazionali (ad esempio, REVE richiedeva 194,79 GFLOPs/s), mantenendo o superando la loro accuratezza.
• Studi di Ablazione:
  • Stato Persistente: L'uso di uno stato nascosto persistente (rispetto al reset ogni 5 secondi) ha migliorato l'AUROC di +0,015 su CHB-MIT e ha migliorato significativamente la localizzazione dell'esordio della crisi.
  • Fasi di Addestramento: Il pre-addestramento completo a due fasi (Fase 1 + Fase 2) ha prodotto guadagni sostanzialmente maggiori (+2,30 AUROC) rispetto al fine-tuning standard o alla sola Fase 1, confermando la necessità di una rifinitura a livello di rappresentazione per i compiti di streaming.

Significato e Affermazioni

Il documento afferma che CaMBRAIN rappresenta un cambiamento di paradigma nella modellazione EEG allineando sia l'architettura che gli obiettivi di addestramento alla natura causale e continua intrinseca dei segnali EEG.

• Capacità in Tempo Reale: Abilita il primo modello capace di eseguire inferenza continua a lungo raggio su segnali EEG di lunghezza variabile senza la penalità computazionale della scalatura quadratica o la ridondanza delle finestre scorrevoli.
• Rilevanza Clinica: L'approccio riflette il modo in cui i clinici monitorano l'EEG al capezzale, mantenendo una memoria compatta dell'intera storia del segnale per rilevare eventi che possono essere separati da lunghi intervalli.
• Efficienza: Sfruttando la scalatura temporale lineare dei Modelli a Spazio di Stato e un obiettivo di addestramento a livello di rappresentazione, CaMBRAIN raggiunge un'alta velocità di elaborazione e bassa latenza, rendendolo adatto per applicazioni cliniche e indossabili in tempo reale dove le risorse computazionali sono limitate.

Gli autori riconoscono che, sebbene CaMBRAIN stabilisca nuovi benchmark su diversi dataset, rimane competitivo piuttosto che superiore nel rilevamento di EEG anomalo su TUAB rispetto ai modelli fondazionali più forti, suggerendo che la scalatura del corpus di pre-addestramento per corrispondere a dataset più grandi e eterogenei (come REVE) è una direzione futura necessaria.