eeg-fm-masking_jepa_r9cm_L2

Pretrained EEG encoder from the paper What masking geometry works best for EEG foundation models? (arXiv:2609.33487 · 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 JEPA (joint-embedding predictive architecture): a predictor maps the encoder's context to the embeddings an EMA teacher produces for the masked patches (no variance/covariance regulariser).

This is the masking geometry recommended by the paper (r = 9 cm, L = 2).

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 JEPA predictor and the EMA teacher are 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

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_jepa_r9cm_L2"
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_jepa_r9cm_L2"
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:

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.

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:

@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},
}