---
license: cc-by-4.0
pipeline_tag: feature-extraction
tags:
- eeg
- self-supervised-learning
- foundation-model
- pytorch
- masked-autoencoder
- open-eeg-bench
---

# eeg-fm-masking_mae_rone_L16

Pretrained EEG encoder from the paper *What masking geometry works best for EEG foundation models?*
([arXiv:2609.33487](https://arxiv.org/abs/2609.33487) · [website](https://pierregtch.github.io/eeg-fm-masking) · [code](https://github.com/PierreGtch/eeg-fm-masking) · [all 58 models](https://huggingface.co/collections/PierreGtch/eeg-fm-masking-6ab912b6a03bba1348fc7366)).

It is one of **58 encoders trained under an identical recipe** where only the masking
geometry changes: 5 spatial radii × 6 temporal lengths × 2 frameworks
(the r = all, L = 33 cell, which would mask the whole window, does not exist).
This model is a **masked autoencoder (MAE)**: the encoder sees the unmasked patches and a light decoder reconstructs the raw signal of the masked patches.

| | |
|---|---|
| Framework | MAE |
| Mask spatial radius `r` | one channel |
| Mask temporal length `L` | 16 patches |
| Masker parameter `pct_unmasked` | 0.45 |
| Checkpoint | end of epoch 10 of 10 (the one evaluated in the paper); epochs 1–9 in `epoch_01/` … `epoch_09/` |
| Encoder parameters | 12.69 M |
| Training run | [`46gmjq13`](https://wandb.ai/pierregtch/chan-inv-clf/runs/46gmjq13) |

The paper recommends r = 9 cm, L = 2: see [`eeg-fm-masking_mae_r9cm_L2`](https://huggingface.co/PierreGtch/eeg-fm-masking_mae_r9cm_L2) and [`eeg-fm-masking_jepa_r9cm_L2`](https://huggingface.co/PierreGtch/eeg-fm-masking_jepa_r9cm_L2).

## What is in this repo

* `model.safetensors`: the **encoder only** (patch tokeniser `feature_encoder.*` + transformer
  `model.*`), i.e. exactly the tensors loaded for the downstream evaluation of the paper.
  The MAE decoder is
  not included; for JEPA the published weights are the student encoder, as evaluated in the paper.
* `epoch_01/` … `epoch_09/`: `model.safetensors` of the intermediate checkpoints (end of epochs
  1 to 9 of the same run), with the same tensors and key names; they use the same `config.json`.
* `config.json`: the keyword arguments of `ContextualEncoderBenchmarkWrapper` (architecture +
  input scaling). Pass it unchanged as `model_kwargs`.
* `metadata.json`: masking parameters, training-run id, checkpoint epoch/step.

## Input requirements

* **Sampling rate: 200 Hz.** The signal is cut into 1 s patches (200 samples, 20-sample overlap).
* **Units: volts.** The wrapper multiplies by `factor = 1e+06` and applies
  per-window `median_std_clip` scaling (clip at σ = 15) itself;
  do not standardise the data beforehand.
* **Channel positions in metres** (MNE `info["chs"][i]["loc"][:3]`). The model is montage-agnostic:
  any number and set of channels works, as long as every channel has a 3D position.

## Usage

Install the code: `pip install git+https://github.com/PierreGtch/eeg-fm-masking`.

### Loading the model

```python
import json
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from eeg_fm_masking.oeb.wrapper import ContextualEncoderBenchmarkWrapper

repo = "PierreGtch/eeg-fm-masking_mae_rone_L16"
config = json.load(open(hf_hub_download(repo, "config.json")))
model = ContextualEncoderBenchmarkWrapper(
    n_chans=n_chans, n_times=n_times, n_outputs=n_outputs, sfreq=200.0,
    chs_info=chs_info,  # MNE channel info (info["chs"]), positions in metres
    **config,
)
model.load_state_dict(load_file(hf_hub_download(repo, "model.safetensors")), strict=False)
```

`strict=False` only leaves out the dataset-dependent parts (channel-position buffer and
classification head).

### Evaluation / fine-tuning with OpenEEGBench

```python
import json
from huggingface_hub import hf_hub_download
from open_eeg_bench.backbone import PretrainedBackbone

repo = "PierreGtch/eeg-fm-masking_mae_rone_L16"
backbone = PretrainedBackbone(
    model_cls="eeg_fm_masking.oeb.wrapper.ContextualEncoderBenchmarkWrapper",
    hub_repo=repo,
    model_kwargs=json.load(open(hf_hub_download(repo, "config.json"))),
)
```

### Intermediate checkpoints (epochs 1–9)

The end-of-epoch checkpoints of the same run are in the subfolders `epoch_01/` … `epoch_09/`
(the final epoch 10 is the `model.safetensors` at the root). Build `model` as above, then:

```python
epoch = 5  # 1 to 9
weights = load_file(hf_hub_download(repo, "model.safetensors", subfolder=f"epoch_{epoch:02d}"))
model.load_state_dict(weights, strict=False)
```

## Downstream results (OpenEEGBench, frozen encoder + ridge probe)

Frozen encoder, ridge regression/classification on the flattened contextual features,
12 datasets × 5 seeds. Balanced accuracy for classification, R² for `seed-vig`.

| Dataset | Metric | Score (mean ± sd) | Seeds |
|---|---|---|---|
| arithmetic_zyma2019 | balanced acc. | 0.710 ± 0.026 | 5 |
| bcic2020-3 | balanced acc. | 0.265 ± 0.015 | 5 |
| bcic2a | balanced acc. | 0.429 ± 0.014 | 5 |
| chbmit | balanced acc. | 0.895 ± 0.023 | 5 |
| faced | balanced acc. | 0.296 ± 0.003 | 5 |
| isruc-sleep | balanced acc. | 0.641 ± 0.003 | 5 |
| mdd_mumtaz2016 | balanced acc. | 0.825 ± 0.009 | 5 |
| physionet | balanced acc. | 0.571 ± 0.004 | 5 |
| seed-v | balanced acc. | 0.285 ± 0.002 | 5 |
| seed-vig | R² | -0.143 ± 0.010 | 5 |
| tuab | balanced acc. | 0.800 ± 0.004 | 5 |
| tuev | balanced acc. | 0.940 ± 0.026 | 5 |

## Training

* **Data:** the openly-licensed subset of the REVE pre-training corpus (323 recordings),
  so that the weights can be redistributed.
* **Schedule:** 10 epochs, 2 × H100, batch size 600 per GPU,
  learning rate 0.00024 (warm-up 3080 steps, final 1e-06),
  weight decay 0.01.

## License and citation

Weights released under **CC-BY-4.0**; code under MIT ([GitHub](https://github.com/PierreGtch/eeg-fm-masking)).
If you use these models, please cite the paper:

```bibtex
@misc{guetschel2026maskinggeometryworksbest,
      title={What masking geometry works best for EEG foundation models?},
      author={Pierre Guetschel and Bruno Aristimunha and Yassine El Ouahidi and Arnaud Delorme and Thomas Moreau and Michael Tangermann},
      year={2026},
      eprint={2609.33487},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2609.33487},
}
```