PierreGtch / eeg-fm-masking_mae_rone_L8

eeg-fm-masking_mae_rone_L8

Pretrained EEG encoder from the paper What masking geometry works best for EEG foundation models? A controlled evaluation across MAE and JEPA ( website · code · all 58 models ).

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.

The paper recommends r = 9 cm, L = 2: see eeg-fm-masking_mae_r9cm_L2 and 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.
• 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

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_L8"
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

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

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

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 .

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 ). If you use these models, please cite the paper (reference on the GitHub page ).