LUNA: Efficient and Topology-Agnostic Foundation Model for EEG Signal Analysis

AI summary

Overview

Research area: Machine learning for biomedical signal processing — specifically self-supervised foundation models for electroencephalography (EEG).

Technical level: Intermediate. The paper assumes familiarity with transformer attention, latent bottlenecks, and self-supervised pretraining objectives, though its central idea is conceptually simple.

Scope: The paper introduces LUNA, a self-supervised EEG foundation model whose channel-unification module makes it topology-agnostic and computationally linear in electrode count, and evaluates it on four downstream EEG benchmarks against supervised and self-supervised baselines.

What This Paper Is About

Public EEG datasets each define their own electrode layout, meaning a model trained on one dataset often degrades sharply when applied to another — the paper cites up to 14 percentage points of accuracy loss for motor-imagery decoders transferring from PhysioNet to KU, and 13–15 percentage-point drops for emotion-recognition models such as BIOT and MMM between SEED and DEAP montages. Existing fixes either train separate models per montage, throw away up to 80% of data by keeping only shared electrodes, or flatten channels and time into one sequence, which costs 𝒪((S·C)²) attention and exhausts memory on dense caps. LUNA's goal is a single architecture that accepts arbitrary electrode configurations and whose compute does not blow up as electrode count grows.

Key Contributions

1. Topology-invariant encoder. A learned-query / cross-attention module that projects arbitrary-sized channel sets into a fixed latent space (ℝ^(Q×E)), a projection that is permutation-invariant with respect to input channels.
2. Linear-in-channels complexity. Patch-wise temporal attention operates only on the fixed latent sequence, reducing the temporal sequence length from S·C to S and decoupling FLOPs and memory from electrode count.
3. State-of-the-art accuracy-efficiency trade-off. The paper reports 81.57% balanced accuracy on TUAB and 39.18% on SEED-V, plus AUROC of 0.921 on TUAR and 0.802 on TUSL, while reducing FLOPs by 300× and GPU memory footprint by up to 10× on high-density recordings, with gains said to hold across diverse electrode configurations.
4. A query specialization loss. An auxiliary objective that penalizes similarity in query–channel affinity matrices, encouraging a diverse, non-redundant set of spatial filters.

Main Findings

• TUAB (abnormal EEG detection): LUNA-Huge (311.4M parameters) reaches 81.57% ± 0.11 balanced accuracy, 0.9029 ± 0.0014 AUC-PR, and 0.8957 ± 0.0011 AUROC. This surpasses most self-supervised baselines and approaches large-scale models such as LaBraM-Huge (82.58% balanced accuracy, 0.9162 AUROC) and CBraMod (82.49%, 0.9156 AUROC), but does not exceed them.
• TUAR (artifact detection): LUNA-Huge achieves 0.921 ± 0.011 AUROC and 0.528 ± 0.012 AUC-PR, described as state-of-the-art and outperforming FEMBA-Large (0.915 AUROC).
• TUSL (slowing event classification): LUNA-Huge reaches 0.802 ± 0.005 AUROC, the highest among all compared models, with AUC-PR of 0.289 ± 0.008.
• SEED-V (emotion recognition, unseen 62-channel topology): LUNA-Large reaches 0.3918 ± 0.0066 balanced accuracy, 0.2073 ± 0.0045 Cohen's Kappa, and 0.3586 ± 0.0013 weighted F1. The paper states this lags leading methods such as CBraMod (0.4091 balanced accuracy) by 2–3 percentage points, which it attributes to reliance on positional encodings learned during pre-training.
• Compute scaling: LUNA maintains near-constant compute cost as channel count increases, in contrast to CBraMod's 𝒪(C²) scaling, and supports thousands of temporal patches without the quadratic cost faced by LaBraM.
• Ablations: Removing frequency features causes the largest drop (up to −0.012 AUROC). Replacing learned queries with fixed spatial regions changes AUROC by only −0.004 to −0.006, and removing the specialization loss changes it by −0.003 to −0.006 — both within seed variation, so the paper emphasizes the practical rather than statistically significant advantage of learned queries.
• Latent structure before fine-tuning: t-SNE visualizations show normal and abnormal EEGs separating into clusters on TUAB and artifact classes partially separating on TUAR even before fine-tuning.
• Emergent query specialization: Learned queries show distinct spatial profiles, some localized (e.g., frontal regions) and others aggregating broader signals.

Methodology in Plain English

LUNA uses an encoder-decoder design. Raw EEG (batch B, channels C, time T) is split into non-overlapping temporal patches of 40 timestamps each. Each patch is embedded two ways: a 1D convolutional network captures temporal features, and an MLP projects the magnitude and phase of the patch's Fourier transform. These are summed. Electrode positions are encoded by applying NeRF-inspired sinusoidal encoding to normalized 3D electrode coordinates, then passing them through an MLP and adding the result to the patch features.

The key step is channel unification. A small set of Q learned queries (Q much smaller than C) cross-attends to the channel features within each patch, producing a fixed-size latent representation. Because attention over a set is permutation-invariant, this step does not care about electrode ordering or count, and its cost grows linearly with channels. Transformer encoder layers then process this latent representation, and a patch-wise temporal encoder with Rotary Positional Embeddings models temporal dependencies across the S patches — a much shorter sequence than S·C.

Pretraining uses a masked-patch reconstruction objective: random tokens are masked, and C learned decoder queries — indexed by channel labels, so they can be reused across datasets where electrodes overlap — cross-attend to the encoder output and a linear projection recovers the patch values. A Smooth L1 loss (β = 1) is applied to masked patches and, with a weight α, to visible patches too. An auxiliary query specialization loss penalizes off-diagonal similarity in query–channel affinity matrices. For fine-tuning, a single aggregation query attends to the encoder output and the pooled result goes to an MLP classifier.

Pretraining data is TUEG plus the Siena Scalp EEG Database, spanning 20, 22, and 29 channels and over 21,000 hours (over 21,900 hours stated in the experimental setup). All downstream subjects and recordings were excluded from pretraining. Preprocessing is minimal: 0.1–75 Hz bandpass, a 50 Hz or 60 Hz notch filter, resampling to 256 Hz, z-score normalization per channel per sample, and a bipolar "double-banana" montage for TUEG/TUAB/TUAR/TUSL while Siena and SEED-V use unipolar format. Models come in Base (7M), Large (43M), and Huge (311M) sizes. Experiments ran on eight NVIDIA A100 GPUs with Python 3.11.6, PyTorch 2.4.1, CUDA 12.1, and bf16 mixed precision; results are reported as mean and standard deviation over three random seeds.

Why This Matters

Impact on research. EEG datasets are fragmented because every lab uses a different cap. A model that accepts arbitrary montages without retraining removes a major barrier to aggregating data across corpora, and the linear-in-channels scaling makes dense, high-channel-count recordings tractable where quadratic spatial attention is impractical. The paper explicitly frames topology-agnostic latent representations as a path toward scalable EEG modeling.

Real-world applications:

• Clinical EEG triage, where abnormality detection (TUAB) and slowing-event classification (TUSL) support neurological diagnosis.
• Automated artifact rejection (TUAR) to clean recordings before expert review, reducing manual annotation load.
• Emotion recognition from brain signals (SEED-V), relevant to human-computer interaction and affective computing.
• Deployment in resource-constrained settings, since the reported FLOPs and memory reductions matter for long recordings and high-density caps.

Industry relevance. The efficiency claims — 300× FLOPs reduction and up to 10× less GPU memory — directly affect inference cost and hardware requirements. The paper's releases include a public code repository (github.com/pulp-bio/biofoundation), and the work was supported by the ETH Future Computing Laboratory, financed by a donation from Huawei Technologies, alongside Swiss National Science Foundation grant 193813 and a Swiss National Supercomputing Centre allocation.

Future Directions

• Improving generalization to unseen topologies. SEED-V performance suggests sensitivity to layouts far from the pretraining montages, likely due to positional encoding constraints; the authors propose enhanced spatial generalization strategies or hybrid learned/geometric embeddings.
• Broadening pretraining montage diversity. Pretraining covers only three dominant layouts; the authors plan multi-dataset pretraining and randomized channel dropout to reach unseen dense and sparse montages.
• Cross-modality unification. Extending the approach to invasive modalities such as sEEG and ECoG.
• Domain priors and real-time use. Integrating neurophysiological constraints and adapting LUNA for real-time inference. The authors also flag ethical concerns — algorithmic bias and patient data privacy — as necessary companions to technical improvements.

Target Audience

Researchers and engineers working on EEG or biosignal foundation models, self-supervised representation learning, and efficient transformer architectures. It is also relevant to clinical-neurophysiology groups interested in cross-dataset transfer, and to practitioners evaluating whether an existing large model can be replaced by a smaller, montage-agnostic one — though readers should note that the paper does not claim an accuracy win over the strongest baselines (LaBraM and CBraMod on TUAB, CBraMod on SEED-V), arguing instead for the accuracy-efficiency trade-off.

Authors’ abstract

Electroencephalography (EEG) offers a non-invasive lens into human brain activity, but building large-scale models is hampered by topological heterogeneity: each public EEG data defines its own electrode layout, limiting generalization. We introduce LUNA (Latent Unified Network Architecture), a self-supervised foundation model that reconciles disparate electrode geometries while scaling linearly -- not quadratically -- with channel count. LUNA compresses multi-channel EEG into a fixed-size, topology-agnostic latent space via learned queries and cross-attention. Downstream transformer blocks then operate exclusively on this latent representation using patch-wise temporal self-attention, decoupling computation from electrode count. Pre-trained on TUEG and Siena (over 21,000 hours of raw EEG across diverse montages) using a masked-patch reconstruction objective, LUNA transfers effectively to four downstream tasks: abnormality detection, artifact rejection, slowing classification, and emotion recognition. It demonstrates highly competitive performance across several benchmarks, achieving state-of-the-art results on TUAR and TUSL, e.g., 0.921 AUROC on TUAR, while reducing FLOPs by 300x and trimming GPU memory use by up to 10x. Critically, these gains are consistent across all evaluated electrode configurations. Code is available at https://github.com/pulp-bio/BioFoundation