---
license: apache-2.0
tags:
- eeg
- neuroscience
- foundation-model
- embeddings
- matryoshka
pipeline_tag: feature-extraction
library_name: neuroencoder
extra_gated_prompt: |-
  The MRL model is currently gated. Access is granted to verified researchers.
  Please briefly describe your institution, role, and intended use.
  If you have a private invitation code, paste it in the "Intended use" field.
extra_gated_fields:
  Institution: text
  Role: text
  Intended use: text
  I agree to use this model for research purposes only: checkbox
---

# EPI Embedding

![](umap.png)

EEG model embeddings, distilled from EPI-250k (trained on ~250,000 hours of clinical EEG).

The model produces a 768-dimensional embedding that you can truncate to **768, 384, 192, 48, or 16** dimensions via [Matryoshka Representation Learning](https://arxiv.org/abs/2205.13147).

## Usage

Install:

```bash
pip install neuroencoder
```

Then:

```python
import mne, neuroencoder as ne
from neuroencoder import MRL

raw = mne.io.read_raw_edf("recording.edf", preload=True)
model = MRL.from_pretrained()                         # auto-downloads on first use

embeddings = model.embed(
    raw.get_data(),
    sfreq=raw.info["sfreq"],
    channel_names=raw.ch_names,
    dim=192,
)
# -> numpy array, shape [N, 192], L2-normalized

ne.explore(embeddings)                                # interactive Apple Embedding Atlas
```

`model.embed` runs the full pipeline (filter -> resample -> 8-region average -> 30s sliding window -> embed) and returns numpy. For more control, split into:

```python
images = ne.preprocess(eeg, sfreq=256, channel_names=ch_names)   # [N, 8, 224, 224]
embeddings = model.predict(images, dim=192)                       # torch tensor on model device
```

## Loading directly from a checkpoint

```python
model = MRL.from_checkpoint("path/to/last.ckpt")
```

Handles both raw state dicts and PyTorch Lightning checkpoint formats.

## Benchmarks

Frozen linear probes, 5-fold subject-level cross-validation. Balanced accuracy (%). The first column is **EPI-250k**, our base foundation model (not publicly released) — the upper bound on what the distilled MRL model can preserve. The remaining columns are the MRL model at each truncation dimension.

### Private clinical tasks

40,909 annotated 30-second epochs from the Swiss Epilepsy Center.

| Task | EPI-250k | 768 | 384 | 192 | 48 | 16 |
|------|:-------:|:---:|:---:|:---:|:--:|:--:|
| Seizure / Wake | **93.4** | 93.1 | 92.7 | 92.5 | 91.5 | 84.1 |
| Sleep (5-class) | **85.1** | 77.0 | 77.4 | 76.9 | 76.5 | 73.2 |
| Artifact / Wake | **90.2** | 90.5 | 90.3 | 90.5 | 90.7 | 65.9 |
| Seizure / Sleep | **88.8** | 85.2 | 84.9 | 84.0 | 82.1 | 79.4 |
| Spike / Seizure | **81.5** | 76.2 | 75.9 | 74.7 | 71.0 | 65.5 |
| Spike / Wake | **97.0** | 94.8 | 94.7 | 94.6 | 92.9 | 87.2 |
| Artifact / Spike | **78.8** | 76.0 | 75.6 | 75.3 | 74.4 | 70.4 |
| Category (6-cls) | **36.3** | 33.6 | 33.3 | 32.8 | 31.7 | 27.4 |
| Clinical Sub (7-cls) | **42.7** | 31.4 | 31.4 | 31.4 | 27.0 | 23.7 |
| All Sublabels (49-cls) | **22.1** | 14.8 | 14.4 | 13.7 | 12.3 | 10.6 |

### Public benchmarks

10 standard public EEG datasets, evaluated under identical conditions.

| Task | EPI-250k | 768 | 384 | 192 | 48 | 16 |
|------|:-------:|:---:|:---:|:---:|:--:|:--:|
| TUAB | **73.1** | 72.4 | 72.5 | 72.9 | 72.2 | 70.4 |
| TUEV | **54.5** | 45.9 | 47.2 | 46.7 | 42.8 | 32.1 |
| TUAR | **45.2** | 43.0 | 42.9 | 42.2 | 39.5 | 36.5 |
| TUSL | **73.3** | 71.5 | 75.1 | 77.1 | 71.3 | 69.7 |
| Mumtaz | **82.1** | 80.7 | 81.8 | 82.6 | 83.2 | 83.1 |
| Schizo | **71.1** | 70.1 | 69.4 | 69.5 | 69.4 | 66.7 |
| MentArith | **60.9** | 60.2 | 59.9 | 58.6 | 55.6 | 52.2 |
| ADFTD | **43.2** | 40.0 | 40.0 | 41.0 | 38.6 | 35.9 |
| PhysioMI | **30.3** | 28.3 | 28.4 | 27.3 | 27.7 | 25.2 |
| Parkinsons | **62.9** | 58.9 | 58.6 | 58.2 | 55.9 | 53.2 |

Numeric column headers (`768`, `384`, ...) are the MRL truncation dimensions.

## Documentation

- Docs: [docs.neuroencoder.com](https://docs.neuroencoder.com)
- GitHub: [github.com/avocardio/neuroencoder](https://github.com/avocardio/neuroencoder)

## Citation

Paper in preparation. A citation will be added once published.