LuMamba: Latent Unified Mamba for Electrode Topology-Invariant and Efficient EEG Modeling

Technical Deep Dive: LuMamba

Overview

LuMamba is a self-supervised framework for efficient and topology-invariant EEG modeling. It combines LUNA's channel-unifying cross-attention mechanism with FEMBA's linear-complexity Mamba state-space modeling backbone, enabling effective EEG representation learning across heterogeneous electrode configurations.

The key innovations of LuMamba are:

• Topology-Invariant SSM for EEG: Fusing LUNA's cross-attention with FEMBA's bi-Mamba blocks allows LuMamba to handle varying electrode counts and placements.
• First Adaptation of LeJEPA to Biosignals: LuMamba adapts the LeJEPA framework, which regularizes embeddings toward an isotropic Gaussian distribution, to EEG time series.
• SSL Objective Trade-offs: LuMamba compares reconstruction-only, LeJEPA-only, and mixed pre-training objectives, revealing complementary benefits for latent structure and downstream generalization.
• Efficiency and Scalability: With only 4.6M parameters, LuMamba requires 26× fewer FLOPs than LUNA and 377× fewer than LaBraM at equivalent sequence lengths, scaling to 12× longer sequences before reaching memory limits.

Problem & Context

Electroencephalography (EEG) is a crucial tool for clinical diagnostics, cognitive neuroscience, and brain-computer interfaces. Recent advances in self-supervised learning have transformed EEG analysis, but a key challenge persists: topological heterogeneity. EEG datasets exhibit wide variation in electrode count and placement, causing severe performance degradation when models are evaluated on unseen configurations.

Prior approaches either train separate models per configuration or discard substantial data by retaining only shared electrodes. LUNA addressed this by learning to project electrode layouts into a fixed latent space, but its Transformer-based architecture incurs high computational cost. Whether such topology-invariant encoding remains effective when combined with efficient state-space models (SSMs) like Mamba is an open question.

Additionally, the optimal self-supervised pre-training objective for SSM architectures is unclear. Current EEG foundation models primarily use masked reconstruction or contrastive learning, but the recently proposed LeJEPA framework, which regularizes embeddings toward an isotropic Gaussian distribution, has not been applied to biosignal time series.

Methodology

Dataset

LuMamba is pre-trained on the Temple University Hospital EEG (TUEG) corpus, a large dataset of over 21,600 hours of recordings from 14,000+ patients. Downstream experiments use task-specific subsets from the TUH family, as well as two datasets with unseen channel montages: APAVA for Alzheimer's disease detection and TDBrain for Parkinson's disease detection.

Architecture

LuMamba's architecture combines several key components:

1. Encoder (LUNA): Tokenizes the input EEG signal and projects each patch into a latent space using 1D convolutions, FFT features, and electrode position encodings.
2. Channel Unification (LUNA): Applies cross-attention to project the channel dimension into a fixed-size latent representation, decoupling the encoder from the input electrode configuration.
3. Temporal Modeling (FEMBA): Processes the latent representation with two bi-Mamba blocks, capturing bidirectional temporal dependencies efficiently.
4. Decoder and Classifier: For pre-training, a cross-attention decoder reconstructs the original channel space. For fine-tuning, a lightweight Mamba-based classification head is used.

Pre-training Objectives

LuMamba systematically investigates the interplay between two self-supervised objectives:

1. Masked Reconstruction: Random input patches are masked for the model to reconstruct, as in prior EEG foundation models.
2. LeJEPA: Adapts the LeJEPA framework to EEG by constructing local and global temporal views and using the Epps-Pulley test to regularize the latent space toward an isotropic Gaussian distribution.

LuMamba evaluates reconstruction-only, LeJEPA-only, and a mixed pre-training strategy that combines both objectives.

Fine-tuning

For downstream tasks, the decoder is replaced with a lightweight Mamba-based classification head. LuMamba is evaluated on five EEG datasets, including three TUH benchmarks (TUAB, TUSL, TUAR) and two datasets with unseen channel montages (TDBrain and APAVA).

Results

Assessing LeJEPA Benefits

• t-SNE visualizations show that reconstruction-only pre-training yields well-separated clusters, while LeJEPA-only produces more diffuse, isotropic embeddings.
• The mixed LeJEPA-reconstruction objective achieves the best of both worlds: strong downstream performance and improved generalization to unseen electrode configurations.
• On APAVA (Alzheimer's detection), the mixed objective improves AUPR by over 20% compared to reconstruction alone.

Comparison to State-of-the-Art

• On the TUAB benchmark, LuMamba achieves performance comparable to LaBraM in Balanced Accuracy and AUROC, with slightly lower AUPR.
• On mental condition tasks with unseen montages, LuMamba demonstrates strong generalization, outperforming prior state-of-the-art models by ~4% AUPR on APAVA (Alzheimer's) and matching performance on TDBrain (Parkinson's).
• LuMamba underperforms task-specific state-of-the-art methods on TUAR and TUSL, which are highly class-imbalanced datasets.

Computational Efficiency and Scalability

• LuMamba consistently requires fewer FLOPs than attention-based foundation models at equal sequence lengths, requiring 26.5× fewer FLOPs than LUNA and 377× fewer than LaBraM.
• LuMamba also exhibits improved scalability, supporting sequences that are 12.6× longer than LUNA and 501× longer than LaBraM before reaching memory limits.

Interpretation

LuMamba successfully combines LUNA's topology-invariant encoding with FEMBA's efficient Mamba blocks, enabling effective EEG representation learning across heterogeneous electrode configurations. The mixed LeJEPA-reconstruction pre-training strategy proves valuable, balancing latent structure and downstream generalization.

LuMamba's performance is competitive with state-of-the-art EEG foundation models on standard benchmarks, while demonstrating strong cross-montage generalization, particularly for disease detection tasks. Its substantial computational efficiency and scalability advantages make it a promising candidate for real-world deployment, especially in resource-constrained settings.

The underperformance on highly imbalanced datasets like TUSL suggests that LuMamba's focus on generalization may come at the cost of specialized performance on certain narrow tasks. Future work could explore techniques to better handle class imbalance while preserving the model's broader capabilities.

Limitations & Uncertainties

• The study is limited to relatively small disease-specific datasets for Alzheimer's and Parkinson's detection. Broader evaluations on larger, more diverse clinical and cognitive neuroscience tasks would further validate the model's generalization.
• The impact of the LeJEPA objective is characterized qualitatively through t-SNE visualizations and quantitatively through downstream task performance. A more detailed theoretical and empirical analysis of the learned representations could provide additional insights.
• While LuMamba demonstrates strong computational efficiency, the study does not explore the model's real-world deployment characteristics, such as inference latency, power consumption, or memory footprint. Practical evaluations in edge computing scenarios would be valuable.

What Comes Next

Future work on LuMamba could explore several promising directions:

1. Scalable Pre-training: Expand the pre-training corpus to further assess the generalizability of the proposed framework, potentially leveraging large-scale web-crawled EEG datasets.
2. Advanced Fine-tuning: Investigate more sophisticated fine-tuning strategies, such as few-shot learning or continual learning, to better handle the unique characteristics of different EEG tasks.
3. Multimodal Integration: Explore ways to integrate LuMamba with other modalities, such as fMRI or MEG, to build truly comprehensive brain-sensing foundation models.
4. Interpretability and Explainability: Develop techniques to better understand the inner workings of LuMamba and the learned representations, potentially leveraging the model's state-space structure.
5. Real-world Deployment: Evaluate LuMamba's performance, efficiency, and robustness in real-world, resource-constrained edge computing scenarios, paving the way for practical clinical and neurotechnology applications.

Sources: [1] LuMamba: Latent Unified Mamba for Electrode Topology-Invariant and Efficient EEG Modeling (arXiv cs.AI preprint)