Upload encoder weights, config and model card

README.md

@@ -0,0 +1,125 @@
+---
+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_r9cm_L1
+
+Pretrained EEG encoder from the paper *What masking geometry works best for EEG foundation models? A controlled evaluation across MAE and JEPA*
+([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` | 9 cm |
+| Mask temporal length `L` | 1 patch |
+| Masker parameter `pct_unmasked` | 0.45 |
+| Checkpoint | epoch 10 of 10 (`v9`, the one evaluated in the paper) |
+| Encoder parameters | 12.69 M |
+| Training run | [`yncl6get`](https://wandb.ai/pierregtch/chan-inv-clf/runs/yncl6get) |
+
+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.
+* `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_r9cm_L1"
+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_r9cm_L1"
+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`.
+
+| Dataset | Metric | Score (mean ± sd) | Seeds |
+|---|---|---|---|
+| arithmetic_zyma2019 | balanced acc. | 0.699 ± 0.011 | 5 |
+| bcic2020-3 | balanced acc. | 0.267 ± 0.029 | 5 |
+| bcic2a | balanced acc. | 0.453 ± 0.005 | 5 |
+| chbmit | balanced acc. | 0.876 ± 0.050 | 5 |
+| faced | balanced acc. | 0.312 ± 0.007 | 5 |
+| isruc-sleep | balanced acc. | 0.662 ± 0.004 | 5 |
+| mdd_mumtaz2016 | balanced acc. | 0.816 ± 0.012 | 5 |
+| physionet | balanced acc. | 0.579 ± 0.004 | 5 |
+| seed-v | balanced acc. | 0.283 ± 0.001 | 5 |
+| seed-vig | R² | -0.163 ± 0.093 | 5 |
+| tuab | balanced acc. | 0.805 ± 0.003 | 5 |
+| tuev | balanced acc. | 0.906 ± 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 (reference on the [GitHub page](https://github.com/PierreGtch/eeg-fm-masking)).

config.json

@@ -0,0 +1,39 @@
+{
+  "clip_sigma": 15.0,
+  "factor": 1000000.0,
+  "feature_encoder": {
+    "dim": 512,
+    "modelName": "LinearPatchEmbedding",
+    "patch_overlap": 20,
+    "patch_size": 200
+  },
+  "masker": {
+    "length_blocks": 1,
+    "n_target_blocks": null,
+    "pct_unmasked": 0.44999999999999996,
+    "radius_blocks": 0.09,
+    "scalp_surface": 0.0942477796076938,
+    "vectorized": true
+  },
+  "pos_encoder": {
+    "max_seconds": 600.0,
+    "max_x": 0.15,
+    "modelName": "AdditivePositionalEncoder",
+    "sfreq_features": 1.1111111111111112,
+    "spat_dim": 384,
+    "time_dim": 128
+  },
+  "scaler": "median_std_clip",
+  "shared_feature_encoder": true,
+  "transformer": {
+    "activation": "gelu",
+    "bias": false,
+    "d_model": 512,
+    "dim_feedforward": 1365,
+    "dropout": 0.0,
+    "glu": true,
+    "nhead": 8,
+    "norm": "rms_norm",
+    "num_layers": 4
+  }
+}

metadata.json

@@ -0,0 +1,23 @@
+{
+  "framework": "mae",
+  "mask_radius_m": 0.09,
+  "mask_radius_label": "9cm",
+  "mask_length_patches": 1,
+  "pct_unmasked": 0.45,
+  "wandb_run_id": "yncl6get",
+  "checkpoint_version": "v9",
+  "wandb_run_name": "mask_sweep_full_mae_r009_l01_af69b36",
+  "wandb_run_url": "https://wandb.ai/pierregtch/chan-inv-clf/runs/yncl6get",
+  "epoch": 9,
+  "global_step": 30064,
+  "n_params": 12687872,
+  "dropped_prefixes": [
+    "predictor"
+  ],
+  "features_shape": [
+    2,
+    19,
+    11,
+    512
+  ]
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0bc057558486b44ee63407ab7317b533d0511f7df5217560058c0b9778c0d9a1
+size 50771472