README.md

---
license: cc-by-nd-4.0
language:
- en
tags:
- eeg
- time-series
- cross-attention
- foundation-model
- neuroscience
library_name: pytorch
---


<div align="center">
  <img src="https://raw.githubusercontent.com/pulp-bio/BioFoundation/refs/heads/main/docs/model/logo/LUNA_logo.png" alt="LUNA Logo" width="800"/>
  <h1>LUNA: Efficient and Topology-Agnostic Foundation Model for EEG</h1>
</div>
<p align="center">
  <a href="https://github.com/pulp-bio/BioFoundation">
    <img src ="https://img.shields.io/github/stars/pulp-bio/BioFoundation?color=ccf" alt="Github">
  </a>
  <a href="https://creativecommons.org/licenses/by-nd/4.0/">
    <img src="https://img.shields.io/badge/License-CC_BY--ND_4.0-lightgrey.svg" alt="License">
  </a>
  <a href="https://arxiv.org/abs/2510.22257">
    <img src="https://img.shields.io/badge/arXiv-2510.22257-b31b1b.svg" alt="Paper">
  </a>
</p>

**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](https://github.com/pulp-bio/BioFoundation) 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):

```yaml
# 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):

```bash
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](https://github.com/pulp-bio/BioFoundation), 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](https://github.com/pulp-bio/BioFoundation).  
- Read `make_datasets/README.md` for exact dataset preparation details.  
- Download the **raw TUSL** dataset from the official [TUH EEG corpus source](https://isip.piconepress.com/projects/nedc/html/tuh_eeg/index.shtml) and place it locally, e.g.: `/eeg_data/TUSL/`.

**1) Prepare data**

```bash
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**

```python
# 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`**

```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**

```bash
python -u run_train.py +experiment=LUNA_finetune
```

*Tip*: to switch sizes without editing the file:

```bash
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:

```bibtex
@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.