LUNA: Efficient and Topology-Agnostic Foundation Model for EEG

LUNA (Latent Unified Network Architecture) is a self-supervised foundation model for EEG that makes models agnostic to electrode topology. LUNA projects arbitrary channel layouts into a fixed-size latent space with learned queries + cross-attention, then runs patch-wise temporal self-attention only on this compact latent. This decouples compute from channel count, yielding linear-in-channels scaling, large FLOPs/memory savings, and strong transfer across datasets and montages.

---

🔒 License & Usage Policy (Weights)

Weights license: The released model weights are licensed under Creative Commons Attribution–NoDerivatives 4.0 (CC BY-ND 4.0). This section summarizes the practical implications for users. This is not legal advice; please read the full license text.

✅ You may

• Use and redistribute the unmodified LUNA weights (including in commercial settings) with proper attribution to the LUNA authors.
• Fine-tune / adapt the weights for your internal use (research or production) without redistributing the modified weights.
• Publish your code, configs, logs, and papers describing experiments with LUNA (please cite the paper).

🚫 You may not

• Share, host, or redistribute any modified weights (including LoRA/adapter/delta checkpoints or pruned/quantized variants). Any parameter set that encodes an adaptation is considered a derivative and cannot be shared under CC BY-ND 4.0.
• Imply endorsement by the LUNA authors for any derivative or evaluation without our written permission.
• Use the LUNA name in a way that suggests your modified model is an official LUNA release.

🤝 How to contribute improvements (PR-gated releases)

We welcome community improvements via a pull-request (PR) workflow. If you believe your improvements should become an official LUNA release:

1. Open a PR in the BioFoundation repository describing the change (architecture/head/training recipe, datasets, preprocessing, compute).
2. Include reproducibility artifacts: configs, seeds, scripts, environment details, training/validation logs, and the evaluation protocol (e.g., TUAB/TUAR/TUSL) with exact splits.
3. Provide comprehensive results (AUROC/AUPR/BA, FLOPs, memory) vs. the baselines reported in the LUNA paper.
4. After maintainer review, approved changes will be retrained/validated and, if accepted, released by the maintainers as a new official LUNA checkpoint under CC BY-ND 4.0.

Rationale: CC BY-ND protects users from fragmented, lower-quality “LUNA variants,” while still enabling internal fine-tuning and a path for the community to upstream improvements through review.

---

🔎 Model Summary

• Goal: Topology-agnostic EEG modeling with linear-in-channels compute/memory.
• Core idea: Channel-Unification Module uses learned queries (Q) with cross-attention to map any set of channels to a fixed latent; temporal Transformer then operates on that latent sequence.
• Pre-training data: TUEG + Siena, >21,000 hours of raw EEG; downstream subjects removed to avoid leakage.
• Downstream tasks: TUAB (abnormal), TUAR (artifacts), TUSL (slowing), SEED-V (emotion; unseen 62-ch montage).

---

🚀 Model Variants

Variant	Parameters
LUNA-Base	7M
LUNA-Large	43M
LUNA-Huge	311M

Scaling increases depth/width of the temporal encoder and the query/embedding sizes in the unification module.

⚙️ Model size configs (ready-made YAMLs)

Pick a LUNA size by selecting one of the provided model configs:

• config/model/LUNA_base.yaml — Base (≈7M)
• config/model/LUNA_large.yaml — Large (≈43M)
• config/model/LUNA_huge.yaml — Huge (≈311M)

Use it via experiment defaults override (recommended):

# inside config/experiment/LUNA_finetune.yaml
defaults:
  - override /data_module: finetune_data_module   # or subject_independent_data_module
  - override /model: LUNA_base                    # change to LUNA_large or LUNA_huge
  - override /scheduler: cosine
  - override /task: finetune_task_LUNA
  - override /criterion: finetune_criterion

Or from the CLI (no file edits):

python -u run_train.py +experiment=LUNA_finetune /model=LUNA_large

---

📊 Results (Highlights)

• TUAR (artifact detection): AUROC 0.921 (LUNA-Huge).
• TUSL (slowing, 4-class): AUROC 0.802 (LUNA-Huge).
• TUAB (abnormal vs normal): Bal. Acc. 81.57%, AUROC 0.8957 (LUNA-Huge).

Efficiency: Up to 300× fewer FLOPs and ≈10× lower GPU memory vs quadratic spatio-temporal attention on dense caps / long windows, thanks to unifying channels before temporal attention.

---

🧠 Intended Use & Limitations

Intended use. Research on EEG representation learning & classification (abnormality, artifacts, slowing, emotion), especially when montages vary or channel counts are high.

Limitations.

• Not a medical device. Do not use for clinical decisions without proper validation & regulatory clearance.
• Unseen topologies: Zero-shot transfer to very different/dense layouts (e.g., SEED-V) can underperform SOTA despite positive scaling; consider augmenting pre-training montage diversity and spatial encodings.
• Distribution shifts: Performance varies across cohorts, devices, and label protocols; validate locally and consider domain adaptation.

---

🏗️ Architecture & Training

Tokenizer & features. EEG is patch-segmented; temporal features via 1D conv w/ GroupNorm+GELU; frequency features (FFT mag/phase → MLP) are added; 3D electrode coordinates encoded via NeRF-style sinusoids → MLP (positional enc).

Channel-Unification Module. Q learned queries cross-attend to channel-wise patch features to produce a fixed Q×E latent per patch; FFN + Transformer layers refine the query tokens. Complexity is O(Q·C) (linear in channels).

Temporal encoder. Patch-wise Transformer with RoPE operates on the latent sequence (length = #patches), not on channels×patches, reducing sequence length and cost substantially.

Pre-training objective. Masked-patch reconstruction with Smooth-L1; decoder uses channel-indexed queries to reconstruct masked tokens. Query specialization loss encourages diverse query–channel affinities.

---

🔧 Fine-tuning — General Checklist

0. Install & read data prep: clone the BioFoundation repo, set up the environment as described there, then open make_datasets/README.md for dataset-specific notes (naming, expected folder layout, and common pitfalls).
1. Choose model size: set - override /model: {LUNA_base|LUNA_large|LUNA_huge} in your experiment YAML (or /model=... via CLI).
2. Point to weights: set pretrained_safetensors_path: /path/to/LUNA_*.safetensors in the experiment YAML.
3. Pick data module:
   • TUH datasets (TUAB/TUSL/TUAR) → - override /data_module: finetune_data_module and optionally override data_module.train/val/test.hdf5_file paths.
   • Non-TUH (e.g., SEED-V) → - override /data_module: subject_independent_data_module and remove the TUH-specific data_module block.
4. Task settings: set classification_type (bc, mc, mmc, mcc) and model.num_classes to match your downstream task.
5. Env vars: export DATA_PATH (dataset root) and CHECKPOINT_DIR (artifacts).
6. Trainer/optimizer: adjust gpus/devices, batch_size, max_epochs, LR/scheduler if needed.
7. I/O: set io.base_output_path and confirm io.checkpoint_dirpath exists.

---

🧪 Example: Fine-tune on TUSL (end-to-end)

0) Install & acquire data

• Follow the installation instructions in the BioFoundation repository.
• Read make_datasets/README.md for exact dataset preparation details.
• Download the raw TUSL dataset from the official TUH EEG corpus source and place it locally, e.g.: /eeg_data/TUSL/.

1) Prepare data

python make_datasets/process_raw_eeg.py tusl       --root_dir /eeg_data/TUSL/edf       --output_dir /processed_eeg

python make_datasets/make_hdf5.py       --prepath /processed_eeg       --dataset TUSL --remove_pkl

2) Set environment variables

# run_train.py (example)
import os
os.environ["DATA_PATH"] = "/processed_eeg"   # contains TUSL_data/{train,val,test}.h5
os.environ["CHECKPOINT_DIR"] = "/LUNA_runs"  # directory for checkpoints & logs

3) Edit the experiment file: config/experiment/LUNA_finetune.yaml

defaults:
  - override /data_module: finetune_data_module # Change based on dataset, finetune_data_module for TUH and subject_independent_data_module for non-TUH
  - override /model: LUNA_base              # Pick the model size, here base, but also available are LUNA_large / LUNA_huge.

pretrained_safetensors_path: /path/to/LUNA_base.safetensors

classification_type: "mcc" # Set based on what type of classification task (Multiclass Classification (MCC), Binary (BC), etc.)
model:
  num_classes: 4 # Set based on how many classes are in your dataset

# Write paths to preprocessed .h5 TUSL files
data_module:
  train:
    _target_: datasets.tuh_dataset.TUH_Dataset
    hdf5_file: ${env:DATA_PATH}/TUSL_data/train.h5 #Here point to the correct file
    finetune: true
  val:
    _target_: datasets.tuh_dataset.TUH_Dataset
    hdf5_file: ${env:DATA_PATH}/TUSL_data/val.h5 #Here point to the correct file
    finetune: true
  test:
    _target_: datasets.tuh_dataset.TUH_Dataset
    hdf5_file: ${env:DATA_PATH}/TUSL_data/test.h5 #Here point to the correct file
    finetune: true

4) Launch

python -u run_train.py +experiment=LUNA_finetune

Tip: to switch sizes without editing the file:

python -u run_train.py +experiment=LUNA_finetune /model=LUNA_large pretrained_safetensors_path=/path/to/LUNA_large.safetensors

---

⚖️ Responsible AI, Risks & Biases

• Clinical safety: research-only; human oversight required.
• Bias & drift: montage/device/population differences can induce shifts; validate and monitor.
• Artifacts & rare events: robustness varies; use QC and task-appropriate preprocessing.

---

🔗 Sources

• Code: https://github.com/pulp-bio/BioFoundation
• Paper: LUNA: Efficient and Topology-Agnostic Foundation Model for EEG Signal Analysis (arxiv:2510.22257).

---

📜 Citation

If you use LUNA, please cite:

@inproceedings{
  doner2025luna,
  title={{LUNA}: Efficient and Topology-Agnostic Foundation Model for {EEG} Signal Analysis},
  author={Berkay D{\"o}ner and Thorir Mar Ingolfsson and Luca Benini and Yawei Li},
  booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
  year={2025},
  url={https://openreview.net/forum?id=uazfjnFL0G}
}

---

🛠️ Maintenance & Contact

• Issues & support: please open a GitHub issue in the BioFoundation repository.

---

🗒️ Changelog

• v1.0: Initial release of LUNA model card with task-specific checkpoints and instructions.