Résumé Technique : CaMBRAIN - Inférence EEG Continue et Temps Réel avec des Modèles à Espace d'État Causal

Énoncé du Problème

Les signaux d'électroencéphalographie (EEG) sont essentiels pour surveiller l'activité cérébrale dans des contextes cliniques et biomécaniques, mais ils présentent des défis uniques pour l'apprentissage profond en raison de leur longueur illimitée, s'étendant souvent de plusieurs minutes à plusieurs heures. Les approches d'apprentissage profond existantes, en particulier celles basées sur les Transformers, font face à deux obstacles majeurs lorsqu'elles sont appliquées à l'EEG de longue durée :

1. Inefficacité Computationnelle : Les mécanismes d'attention entraînent une mise à l'échelle quadratique ($O(T^2)$) avec la longueur de la séquence, les rendant inefficaces pour les signaux longs.
2. Limitations du Contexte : Pour gérer les coûts computationnels, les modèles actuels traitent généralement l'EEG via des fenêtres glissantes avec des longueurs d'entrée fixes. Cette approche force le modèle à recalculer les caractéristiques sur des régions chevauchantes et empêche le maintien d'une représentation persistante de l'historique complet du signal. Par conséquent, ces modèles manquent de compréhension temporelle globale et ne peuvent pas effectuer une véritable inférence continue et temps réel où chaque nouvelle observation s'appuie sur le contexte précédent sans réexaminer les entrées passées.

De plus, les objectifs de préentraînement auto-supervisé existants pour l'EEG reposent souvent sur la reconstruction du signal (par exemple, les autoencodeurs masqués). Bien qu'efficaces pour apprendre la fidélité locale du signal, ces objectifs n'entraînent pas explicitement l'état caché à retenir un contexte à long terme saillant, ce qui est essentiel pour l'inférence en flux continu où le modèle doit compresser l'historique dans une mémoire compacte.

Méthodologie

Les auteurs proposent CaMBRAIN, le premier modèle EEG basé sur Mamba et causal, conçu pour une inférence en flux continu et temps réel. La méthodologie repose sur deux piliers fondamentaux : une architecture causale et un pipeline d'entraînement auto-supervisé spécialisé en plusieurs étapes.

1. Architecture : Mamba-3 Causal

CaMBRAIN utilise un Modèle à Espace d'État (SSM) causal et unidirectionnel, spécifiquement l'architecture Mamba-3, pour traiter les signaux EEG.

• Traitement de l'Entrée : Les signaux EEG bruts sont partitionnés en patches non chevauchants de 62,5 ms (16 échantillons à 256 Hz) sur tous les canaux. Ces patches sont intégrés en tokens à l'aide d'un intégrateur de canal basé sur un Transformer.
• État Caché Persistant : Contrairement aux approches par fenêtre glissante, CaMBRAIN maintient un état caché persistant qui compresse de manière incrémentielle tout le contexte passé. À mesure que de nouveaux patches arrivent, l'état caché est mis à jour de manière récurrente ($h_t = \bar{A}h_{t-1} + \bar{B}x_t$) sans réexaminer les entrées précédentes.
• Contrainte Causale : Le modèle opère strictement de manière unidirectionnelle, garantissant que les prédictions au temps $t$ dépendent uniquement des entrées passées et actuelles ($x_{\le t}$), s'alignant ainsi sur la réalité physique de la surveillance EEG.
• Transitions à Valeurs Complexes : En tirant parti de Mamba-3, le modèle emploie des transitions d'état à valeurs complexes pour mieux capturer les dynamiques oscillatoires et les relations de phase inhérentes aux signaux EEG multi-bandes.

2. Pipeline d'Entraînement

Pour pallier les limites des objectifs basés sur la reconstruction pour le flux continu, les auteurs introduisent un pipeline d'entraînement en trois étapes :

• Étape 1 : Préentraînement Prédictif Causal :
  L'encodeur est préentraîné à l'aide d'une perte causale hybride pour instaurer une structure temporelle générale. Cela combine :

  • Modélisation Autoregressive (ARM) : Prédire l'intégration du patch suivant à partir de l'état caché actuel.
  • Reconstruction Masquée : Reconstruire les tokens EEG bruts masqués à partir du contexte causal fourni par l'encodeur.
    Cette étape garantit que le modèle apprend les dynamiques temporelles locales et la fidélité du signal.

• Étape 2 : Raffinement de la Représentation Latente sans Reconstruction (style JEPA) :
  Pour encourager la rétention d'informations prédictives à long terme, le modèle subit une phase d'entraînement sans décodeur inspirée des Architectures Prédictives d'Intégration Commune (JEPA).

  • Cadre Élève-Professeur : Un réseau élève traite une vue masquée de l'entrée, tandis qu'un réseau professeur (une Moyenne Mobile Exponentielle de l'élève) traite l'entrée non masquée.
  • Objectifs : L'élève est entraîné à prédire :
    1. Latents Masqués : Les représentations du professeur pour les positions masquées.
    2. Latents Futurs : Représentations futures multi-étapes ($K$ étapes en avant) à partir de l'état caché actuel.
       Cette étape force explicitement l'état caché à encoder des caractéristiques informatives, invariantes et temporellement pertinentes nécessaires à l'inférence continue, plutôt que de simplement reconstruire des détails locaux du signal.

• Étape 3 : Affinage pour Tâches en Aval :
  Le modèle préentraîné est affiné sur des tâches spécifiques (par exemple, détection de crises, classification d'artefacts) en utilisant un apprentissage supervisé standard. Notamment, CaMBRAIN n'utilise que l'état caché final pour la classification, évitant ainsi le besoin de pooling ou d'agrégation de tokens sur des segments fixes.

Contributions Clés

1. Formulation de l'Inférence Causale : L'article propose une formulation d'inférence EEG unidirectionnelle qui s'éloigne des approches par fenêtre glissante et bidirectionnelles, traitant les signaux comme des flux continus plutôt que comme des segments indépendants.
2. Architecture CaMBRAIN : Un encodeur Mamba-3 causal qui remplace le recalcul redondant des fenêtres glissantes par un état caché persistant, permettant une inférence efficace et évolutive sur des signaux arbitrairement longs avec une complexité linéaire dans le temps.
3. Entraînement au Niveau de la Représentation : Un cadre auto-supervisé multi-étapes novateur qui va au-delà de la reconstruction du signal. En combinant la prédiction autoregressive avec la prédiction de latents masqués et futurs, il entraîne explicitement l'état caché à retenir un contexte à long terme saillant.

Résultats

CaMBRAIN a été évalué sur quatre ensembles de données EEG divers : TUAR (détection d'artefacts), TUAB (détection d'EEG anormal), MAT (détection de stress mental) et CHB-MIT (détection de crises).

• Performance : CaMBRAIN a obtenu des résultats à l'état de l'art (SOTA) ou compétitifs sur tous les benchmarks.
  • TUAR : A atteint une AUROC de 0,936, surpassant des modèles de base plus grands comme LUNA-H (0,921).
  • CHB-MIT : A atteint une AUROC de 0,921, démontrant une forte performance sur des données de crises hautement déséquilibrées.
  • MAT : A obtenu la meilleure précision équilibrée (0,778) et AUROC (0,876) parmi les méthodes comparées.
• Efficacité Computationnelle : Le modèle a démontré des gains d'efficacité significatifs dans les scénarios de flux continu.
  • À un taux de mise à jour en flux continu de 16 Hz, CaMBRAIN nécessitait 1,23 GFLOPs/s.
  • Cela représente une réduction de 6,1× à 158× du calcul soutenu par rapport à d'autres modèles de base (par exemple, REVE nécessitait 194,79 GFLOPs/s), tout en égalant ou dépassant leur précision.
• Études d'Ablation :
  • État Persistant : L'utilisation d'un état caché persistant (par opposition à une réinitialisation toutes les 5 secondes) a amélioré l'AUC de +0,015 sur CHB-MIT et a considérablement amélioré la localisation du début des crises.
  • Étapes d'Entraînement : Le préentraînement complet en deux étapes (Étape 1 + Étape 2) a produit des gains substantiellement plus importants (+2,30 AUROC) par rapport à un affinage standard ou à l'Étape 1 seule, confirmant la nécessité d'un raffinement au niveau de la représentation pour les tâches de flux continu.

Importance et Revendications

L'article revendique que CaMBRAIN représente un changement de paradigme dans la modélisation EEG en alignant à la fois l'architecture et les objectifs d'entraînement sur la nature causale et continue inhérente des signaux EEG.

• Capacité Temps Réel : Il permet le premier modèle capable d'effectuer une inférence continue à long terme sur des signaux EEG de longueur variable sans la pénalité computationnelle de la mise à l'échelle quadratique ni la redondance des fenêtres glissantes.
• Pertinence Clinique : L'approche reflète la façon dont les cliniciens surveillent l'EEG au chevet du patient, en maintenant une mémoire compacte de l'historique complet du signal pour détecter des événements qui peuvent être séparés par de longs intervalles.
• Efficacité : En tirant parti de la mise à l'échelle linéaire dans le temps des Modèles à Espace d'État et d'un objectif d'entraînement au niveau de la représentation, CaMBRAIN atteint un débit élevé et une faible latence, le rendant adapté aux applications cliniques et portables en temps réel où les ressources computationnelles sont limitées.

Les auteurs reconnaissent que, bien que CaMBRAIN établisse de nouveaux benchmarks sur plusieurs ensembles de données, il reste compétitif plutôt que supérieur sur la détection d'EEG anormal TUAB par rapport aux modèles de base les plus performants, suggérant que l'augmentation de l'échelle du corpus de préentraînement pour correspondre à des ensembles de données plus grands et plus hétérogènes (comme REVE) est une direction future nécessaire.