FEMBA on the Edge: Physiologically-Aware Pre-Training, Quantization, and Deployment of a Bidirectional Mamba EEG...

Area of Science:

• Neuroscience
• Computer Science
• Biomedical Engineering

Background:

• Transformer models for Electroencephalography (EEG) face computational bottlenecks for wearable devices.
• State-Space Models (SSMs) offer efficiency but present quantization challenges.
• Continuous neuro-monitoring is crucial for diagnosing and managing neurological conditions like epilepsy and sleep disorders.

Purpose of the Study:

• To develop a computationally efficient framework for deploying large-scale EEG foundation models on ultra-low-power wearable devices.
• To overcome the limitations of existing models for real-time, continuous neuro-monitoring.
• To enable robust clinical analysis on edge hardware without performance degradation.

Main Methods:

• Introduced FEMBA, a bidirectional Mamba architecture pre-trained on over 21,000 hours of EEG data.
• Developed a novel Physiologically-Aware pre-training objective using low-pass filtering to prioritize neural oscillations.
• Employed Quantization-Aware Training (QAT) to compress the model to 2-bit weights and deployed on a RISC-V microcontroller (GAP9).

Main Results:

• Physiologically-Aware pre-training improved downstream AUROC by 3.5% and AUPR by 3.6% on the TUAB dataset.
• QAT successfully compressed model weights to 2-bit with negligible performance loss, unlike post-training quantization which degraded accuracy by ~30%.
• The embedded implementation achieved real-time inference, reduced memory footprint by 74% to ~2 MB, and used up to 27x fewer FLOPs than Transformer benchmarks.

Conclusions:

• FEMBA demonstrates the feasibility of effectively quantizing and deploying Mamba-based foundation models on extreme-edge hardware for robust clinical analysis.
• This work presents the first full-stack framework for deploying large-scale EEG foundation models on ultra-low-power wearables.
• The developed framework facilitates continuous, SSM-based monitoring for epilepsy and sleep disorders, paving the way for advanced wearable neuro-monitoring solutions.