Technische Samenvatting: CaMBRAIN - Realtime, Continue EEG-Inferentie met Causale State Space-modellen

Probleemstelling

Elektro-encefalografie (EEG)-signalen zijn cruciaal voor het bewaken van hersenactiviteit in klinische en biomechanische omgevingen, maar ze vormen unieke uitdagingen voor deep learning vanwege hun onbeperkte lengte, vaak variërend van minuten tot uren. Bestaande deep learning-benaderingen, met name die gebaseerd op Transformers, staan voor twee hoofdhobbels bij toepassing op EEG met lange duur:

1. Rekenkundige Inefficiëntie: Attention-mechanismen vertonen kwadratische schaling ($O(T^2)$) met de sequentielengte, waardoor ze inefficiënt zijn voor lange signalen.
2. Contextbeperkingen: Om rekenkosten te beheersen, verwerken huidige modellen EEG doorgaans via schuifvensters met vaste invoerlengtes. Deze aanpak dwingt het model om features over overlappende gebieden opnieuw te berekenen en verhindert het behoud van een persistente representatie van het volledige signaalverleden. Bijgevolg missen deze modellen een globaal temporaal begrip en kunnen ze geen echte realtime, continue inferentie uitvoeren waarbij elke nieuwe observatie voortbouwt op eerdere context zonder terug te keren naar eerdere invoer.

Bovendien vertrouwen bestaande zelftoezicht-pretrainingsdoelen voor EEG vaak op signaalreconstructie (bijvoorbeeld masked autoencoders). Hoewel deze effectief zijn voor het leren van lokale signaalfideliteit, trainen ze de verborgen staat niet expliciet om opvallende langetermijncontext te behouden, wat essentieel is voor streaming-inferentie waarbij het model de geschiedenis moet comprimeren tot een compact geheugen.

Methodologie

De auteurs stellen CaMBRAIN voor, het eerste causale, Mamba-gebaseerde EEG-model dat is ontworpen voor realtime, continue streaming-inferentie. De methodologie rust op twee pijlers: een causale architectuur en een gespecialiseerde meerstaps zelftoezicht-trainingspijplijn.

1. Architectuur: Causal Mamba-3

CaMBRAIN maakt gebruik van een unidirectioneel, causaal State Space Model (SSM), specifiek de Mamba-3-architectuur, om EEG-signalen te verwerken.

• Invoerverwerking: Ruwe EEG-signalen worden opgedeeld in niet-overlappende patches van 62,5 ms (16 samples bij 256 Hz) over alle kanalen. Deze patches worden ingebed in tokens met behulp van een transformer-gebaseerde kanaalembedder.
• Persistente Verborgen Staat: In tegenstelling tot schuifvensterbenaderingen, behoudt CaMBRAIN een persistente verborgen staat die alle vorige context stapsgewijs comprimeert. Naarmate nieuwe patches arriveren, wordt de verborgen staat recurrent geüpdatet ($h_t = \bar{A}h_{t-1} + \bar{B}x_t$) zonder terug te keren naar eerdere invoer.
• Causale Beperking: Het model werkt strikt in een alleen-voorwaartse richting, zodat voorspellingen op tijdstip $t$ uitsluitend afhankelijk zijn van vorige en huidige invoer ($x_{\le t}$), in lijn met de fysieke realiteit van EEG-bewaking.
• Complexe Overgangen: Door gebruik te maken van Mamba-3, hanteert het model complexe overgangen in de staat om de oscillatoire dynamiek en faseverhoudingen die inherent zijn aan multi-band EEG-signalen beter vast te leggen.

2. Trainingspijplijn

Om de beperkingen van reconstructie-gebaseerde doelen voor streaming aan te pakken, introduceren de auteurs een drie-staps trainingspijplijn:

• Fase 1: Causale Predictieve Pretraining:
  De encoder wordt voorgeprent met een hybride causale loss om een algemene temporele structuur in te bouwen. Dit combineert:

  • Autoregressieve Modellering (ARM): Het voorspellen van de volgende patch-embedding vanuit de huidige verborgen staat.
  • Gemaskeerde Reconstructie: Het reconstrueren van gemaskeerde ruwe EEG-tokens vanuit de causale context die door de encoder wordt geboden.
    Deze fase zorgt ervoor dat het model lokale temporele dynamiek en signaalfideliteit leert.

• Fase 2: Reconstructievrije Latente Representatieverfijning (JEPA-stijl):
  Om het behoud van langetermijn, voorspellende informatie te stimuleren, ondergaat het model een trainingsfase zonder decoder, geïnspireerd door Joint-Embedding Predictive Architectures (JEPA).

  • Student-Leraar Kader: Een studentnetwerk verwerkt een gemaskeerde weergave van de invoer, terwijl een leraarnetwerk (een Exponentiële Glijdende Gemiddelde van de student) de ongemaskeerde invoer verwerkt.
  • Doelen: De student wordt getraind om te voorspellen:
    1. Gemaskeerde Latenten: De representaties van de leraar voor gemaskeerde posities.
    2. Toekomstige Latenten: Meerdere-staps toekomstige representaties ($K$ stappen vooruit) vanuit de huidige verborgen staat.
       Deze fase dwingt de verborgen staat expliciet om informatieve, invariantie en temporeel relevante features te coderen die noodzakelijk zijn voor continue inferentie, in plaats van alleen lokale signaaldetails te reconstrueren.

• Fase 3: Downstream Fine-tuning:
  Het voorgeprente model wordt gefine-tuned op specifieke taken (bijvoorbeeld epilepsie-aanvalsdetectie, artefactclassificatie) met standaard supervised learning. Opmerkelijk is dat CaMBRAIN uitsluitend de laatste verborgen staat gebruikt voor classificatie, waardoor er geen behoefte is aan pooling of het aggregeren van tokens over vaste segmenten.

Belangrijkste Bijdragen

1. Causale Inferentieformulering: Het artikel stelt een unidirectionele EEG-inferentieformulering voor die afwijkt van schuifvenster- en bidirectionele transformer-benaderingen, door signalen te behandelen als continue streams in plaats van onafhankelijke segmenten.
2. CaMBRAIN-architectuur: Een causale Mamba-3-encoder die redundante schuifvenster-herberekening vervangt door een persistente verborgen staat, waardoor efficiënte, schaalbare inferentie mogelijk wordt over willekeurig lange signalen met lineaire tijdscomplexiteit.
3. Representatieniveau-training: Een nieuw meerstaps zelftoezicht-kader dat verder gaat dan signaalreconstructie. Door autoregressieve voorspelling te combineren met gemaskeerde en toekomstige latente voorspelling, traint het de verborgen staat expliciet om opvallende langetermijncontext te behouden.

Resultaten

CaMBRAIN werd geëvalueerd op vier diverse EEG-datasets: TUAR (artefactdetectie), TUAB (detectie van abnormaal EEG), MAT (detectie van mentale stress) en CHB-MIT (epilepsie-aanvalsdetectie).

• Prestaties: CaMBRAIN behaalde state-of-the-art (SOTA) of concurrerende resultaten op alle benchmarks.
  • TUAR: Behaalde een AUROC van 0,936, wat beter is dan grotere foundation-modellen zoals LUNA-H (0,921).
  • CHB-MIT: Behaalde een AUROC van 0,921, wat sterke prestaties aantoont op sterk onbalansseerde epilepsie-aanvaldata.
  • MAT: Behaalde de beste gebalanceerde nauwkeurigheid (0,778) en AUROC (0,876) onder de vergeleken methoden.
• Rekenkundige Efficiëntie: Het model toonde aanzienlijke winst in efficiëntie in streaming-scenario's.
  • Bij een streaming-updatefrequentie van 16 Hz vereiste CaMBRAIN 1,23 GFLOPs/s.
  • Dit vertegenwoordigt een reductie in doorlopende rekkracht van 6,1× tot 158× ten opzichte van andere foundation-modelbaselines (bijvoorbeeld vereiste REVE 194,79 GFLOPs/s), terwijl de nauwkeurigheid gelijk bleef of werd overtroffen.
• Ablatiestudies:
  • Persistente Staat: Het gebruik van een persistente verborgen staat (in plaats van elke 5 seconden resetten) verbeterde de AUROC met +0,015 op CHB-MIT en verbeterde significant de lokalisatie van het begin van epilepsie-aanvallen.
  • Trainingsfasen: De volledige twee-staps pretraining (Fase 1 + Fase 2) leverde aanzienlijk grotere winsten (+2,30 AUROC) op dan standaard fine-tuning of Fase 1 alleen, wat de noodzaak bevestigt van representatieniveau-verfijning voor streaming-taken.

Betekenis en Claims

Het artikel stelt dat CaMBRAIN een paradigma-verschuiving in EEG-modellering vertegenwoordigt door zowel de architectuur als de trainingsdoelen af te stemmen op de inherente causale en continue aard van EEG-signalen.

• Realtime Capaciteit: Het stelt het eerste model in staat om langetermijn, continue inferentie uit te voeren op EEG-signalen met variabele lengte zonder de rekenkundige straf van kwadratische schaling of de redundantie van schuifvensters.
• Klinische Relevantie: De aanpak weerspiegelt hoe clinici EEG aan het bed bewaken, door een compact geheugen van het volledige signaalverleden te behouden om gebeurtenissen te detecteren die door lange intervallen gescheiden kunnen zijn.
• Efficiëntie: Door gebruik te maken van de lineaire tijdschaling van State Space-modellen en een trainingsdoel op representatieniveau, bereikt CaMBRAIN een hoge doorvoer en lage latentie, waardoor het geschikt is voor realtime klinische en draagbare toepassingen waar rekenbronnen beperkt zijn.

De auteurs erkennen dat hoewel CaMBRAIN nieuwe benchmarks zet op meerdere datasets, het concurrerend blijft in plaats van superieur bij de detectie van abnormaal EEG op TUAB vergeleken met de sterkste foundation-modellen, wat suggereert dat het schalen van het pretraining-corpus om te matchen met grotere, heterogener datasets (zoals REVE) een noodzakelijke toekomstige richting is.