Training a model – masked prediction on EEG

Every track accepts two kinds of entry: a task-specific model, trained on that track’s task alone, and a foundation model, one network reused across tasks rather than rebuilt for each. Neither is privileged, and how you obtain a foundation model is up to you – any data, any objective.

This page is a working template for the foundation-model route: pretrain a small encoder by masked prediction on unlabelled EEG with neuraltrain’s ssl_example project, then score it with neuralbench. Masked prediction is used here only because it needs no labels, so it can pool datasets that share nothing but being EEG. Expect a baseline rather than a competitive entry: the encoder is small and the corpus is four datasets.

Note

Already have a model of your own? Skip to Evaluating a model of your own. Nothing about the competition requires neuraltrain, neuralset or PyTorch Lightning.

1. Install, and download the corpus

ssl_example ships in the repository rather than the wheel, its encoder sits behind the models extra, and Stieger2021Continuous needs moabb:

git clone https://github.com/facebookresearch/neuroai
cd neuroai
pip install './neuraltrain-repo[lightning,models]' 'moabb>=1.7.1'

It pretrains on four EEG datasets – those behind tracks 1-3 (Gifford2022Large, Stieger2021Continuous, Kemp2000Analysis) plus resting-state Miltiadous2023Dice – some 240 subjects in all.

Repoint the paths before the first run. DATADIR, CACHEDIR and SAVEDIR sit at the top of defaults.py, all three under ~/.cache/neuralset and independent of the DATA_DIR you configured for neuralbench. The four datasets want well over 1 TB (Stieger2021Continuous ~940 GB, Gifford2022Large ~220 GB).

Training reads what is on disk and never fetches, so download first – once per machine:

python -m ssl_example.grids.download

Preprocessing is then cached into CACHEDIR on first use and reused by every later run and grid job; changing it pays for that pass again.

2. Check the wiring before paying for it

The debug config swaps the four datasets for MNE’s sample recording, downloaded on first use, and runs a single batch:

cd neuraltrain-repo
python -m ssl_example.grids.test_run

3. Pretrain the encoder

The loop it runs is MAE on EEG: each window is cut into time patches one channel at a time, so a token is one channel over one patch; a random mask_ratio of those tokens is swapped for a learned mask token; and a linear head reconstructs the hidden patches from the encoder’s output, scored on those patches alone. Only the encoder is kept. MAEEG applies that objective to EEG at this scale, and ST-EEGFormer takes it to a foundation model – it won last year’s edition of this challenge, which makes it a useful model for where to take the template below.

The real run finishes on a line reading Pretrained encoder: <SAVEDIR>/ssl_example.main.Experiment.run,1/<uid>/encoder.ckpt. The <uid> is a hash of the config, so copy that path rather than reconstruct it:

python -m ssl_example.grids.defaults

Show ssl_example/grids/defaults.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""Default configuration for MAE pretraining on the challenge datasets."""

from pathlib import Path

import neuralset as ns

PROJECT_NAME = "challenge_mae"
CACHEDIR = f"{ns.CACHE_FOLDER}/cache/{PROJECT_NAME}"
SAVEDIR = f"{ns.CACHE_FOLDER}/results/{PROJECT_NAME}"
# each study claims its own subfolder
DATADIR = f"{ns.CACHE_FOLDER}/data"
for path in [CACHEDIR, SAVEDIR, DATADIR]:
    Path(path).mkdir(parents=True, exist_ok=True)

# preprocessing below matches `neuralbench/defaults/config.yaml`, so the encoder
# sees the same signal downstream as it does here
FREQUENCY = 120.0
WINDOW = 4.0  # 15 patches of 32 samples, hence 15 * n_channels tokens

# tracks 1-3 plus a resting-state study; track 4 is EMG, a different sensor space
STUDIES = [
    "Gifford2022Large",  # track 1, image decoding, 63 ch
    "Stieger2021Continuous",  # track 2, motor imagery, 60 ch
    "Kemp2000Analysis",  # track 3, sleep staging, 2 bipolar derivations
    "Miltiadous2023Dice",  # resting-state, eyes closed, 19 ch
]

default_config = {
    "infra": {
        "cluster": None,  # Run example locally
        "folder": SAVEDIR,
        "gpus_per_node": 1,
        "cpus_per_task": 10,
    },
    "data": {
        "studies": [
            [
                {
                    "name": name,
                    "path": DATADIR,
                    "query": None,
                    "infra": {"backend": "Cached", "folder": CACHEDIR},
                },
                # whole subjects held out, so validation measures generalisation to
                # an unseen recording; pretraining leaves the test split untouched
                {
                    "name": "SklearnSplit",
                    "split_by": "subject",
                    "valid_split_ratio": 0.1,
                    "test_split_ratio": 0.1,
                    "valid_random_state": 33,
                    "test_random_state": 33,
                },
            ]
            for name in STUDIES
        ],
        "segmenter": {
            "extractors": {
                # No "target" extractor: the input is its own target.
                "input": {
                    "name": "EegExtractor",
                    "frequency": FREQUENCY,
                    "filter": (0.1, 75.0),
                    "notch_filter": [50.0, 60.0],
                    "scaler": "RobustScaler",
                    "clamp": 20.0,
                    "infra": {
                        "keep_in_ram": True,
                        "folder": CACHEDIR,
                        "cluster": None,
                    },
                },
            },
            # the recording itself is the trigger: windows tile it, no events needed
            "trigger_query": "type == 'Eeg'",
            "stride": WINDOW,
            "duration": WINDOW,
        },
        # standard_1020 also resolves Sleep-EDF's bipolar names (Fpz-Cz -> Fpz)
        "channel_positions": {
            "n_spatial_dims": 3,
            "layout_or_montage_name": "standard_1020",
            # positions are padded to the channel union of the studies, which is
            # not part of the cache key: kept on disk, they would be served at the
            # width of whichever study set filled the cache first. Recomputing
            # them costs seconds, and `keep_in_ram` still spares repeated reads.
            "infra": {
                "keep_in_ram": True,
                "folder": None,
                "cluster": None,
            },
        },
        "batch_size": 16,
    },
    "brain_model_config": {
        "name": "MaeEncoder",
        "dim": 256,
        "patch_size": 32,
        # `channel_emb_config` left at its default: naming it resets n_dims to 2
    },
    "mask_ratio": 0.5,
    # the module scores the hidden patches only, so a plain MSE is the MAE loss
    "loss": {"name": "MSELoss"},
    "optim": {
        "optimizer": {
            "name": "AdamW",
            "lr": 1e-4,
            "kwargs": {"weight_decay": 0.05},
        },
        "scheduler": {
            "name": "OneCycleLR",
            "kwargs": {"max_lr": 1e-3, "pct_start": 0.2},
        },
    },
    "csv_config": {
        "name": PROJECT_NAME,
        "flush_logs_every_n_steps": 100,
    },
    # set to None to train without Weights & Biases
    "wandb_config": {
        "log_model": False,
        "group": PROJECT_NAME,
        "project": PROJECT_NAME,
    },
    "n_epochs": 50,
    "limit_train_batches": None,
    "patience": 10,
    "fast_dev_run": False,
    "seed": 33,
}


if __name__ == "__main__":
    # The following can be used for local debugging/quick tests.

    from ..main import Experiment

    exp = Experiment(**default_config)

    exp.infra.clear_job()
    out = exp.run()
    print(out)
    print(f"Pretrained encoder: {exp.checkpoint_path}")

Three parts of that config are what make it self-supervised, and the parts to keep when you swap in your own data:

• Windows come from a stride, not from events, so they tile the recording instead of clustering around stimuli.

• There is no target extractor: the input is its own target.

• The split holds out whole subjects, because striding turns one recording into hundreds of near-duplicate windows and splitting over windows would mostly measure memorisation.

mask_ratio is the knob that matters most – hide too little and reconstruction becomes trivial copying. ssl_example/grids/run_grid.py sweeps it on SLURM, and the run works unchanged on several GPUs.

One design choice explains why a single encoder can span a 63-channel cap and Sleep-EDF’s two bipolar derivations, and why adding a dataset below needs no code: a channel is identified by a Fourier embedding of its 3D position on the head, never by its index, and zero-padded channels are dropped from the attention rather than read as signal.

4. Score it on a downstream task

Pretraining is worth only what its representations are worth. neuralbench ships an mae model config that rebuilds this encoder, and --checkpoint points it at your weights:

neuralbench eeg motor_imagery -m mae \
    --checkpoint <the encoder.ckpt path printed above> \
    -w linear_probe_mean

-w linear_probe_mean freezes the encoder and trains only a linear probe on the mean-pooled tokens, at the learning rate the benchmark uses for its own probes – which is what makes the score a property of the pretrained representation, and comparable to the published numbers. Leave -w out to fine-tune end to end; -w lora_r4_flatten sits in between.

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

# Reference self-supervised baseline, pretrained by `neuraltrain-repo/ssl_example`.
# It has no published weights, so use it with `--checkpoint <your encoder>`.
#
# `dim` and `patch_size` must match the pretraining run: on a mismatch
# `load_checkpoint` only logs "Size mismatch" and keeps random weights.
# Channel count need not match, as channels are identified by their position.
#
# `data.neuro` is deliberately absent: `ssl_example` pretrains with the
# preprocessing from `defaults/config.yaml`, so inheriting it keeps the two
# sides in step without restating any of it here.
data:
  channel_positions:
    n_spatial_dims: 3
    # same montage as pretraining; also resolves bipolar derivations
    layout_or_montage_name: standard_1020
brain_model_config:
  =replace=: true
  name: MaeEncoder
  dim: 256
  patch_size: 32
  # `channel_emb_config` absent here and in `ssl_example`, so neither can drift
# n_outputs=None: the probe pools the encoder's tokens; absent channels are zeroed
downstream_model_wrapper:
  model_output_key: null
  aggregation: mean
  probe_config: linear
  # encoder.layers.*.1.{to_q,to_k,to_v,to_out}
  lora_target_modules:
    - to_q
    - to_k
    - to_v
    - to_out

To confirm the pretraining bought you anything, compare against the same architecture with no checkpoint (drop --checkpoint) and against the task-specific baselines on the track pages. The example pretrains on the data of tracks 1-3, so a probe scored there has already seen that data unlabelled – allowed, but silent on generalising to an unseen dataset.

Warning

mae.yaml’s dim and patch_size must match the encoder you pretrained, and nothing checks that they do. On a mismatch neuralbench logs Size mismatch and keeps the randomly initialised layer, which reads as a failed pretraining run rather than a misconfiguration – so check the log before trusting a score. Channel count is the one thing you never have to match.

5. Scale it up

Change the config rather than the code:

• More data: add any study from the NeuralFetch catalog to STUDIES. Unlabelled EEG is the one resource pretraining scales with, so this matters more than any architecture choice. The competition also points at EEGDash for several hundred further EEG corpora.

• A bigger encoder: raise brain_model_config.dim and transformer_config.depth, mirroring any dim or patch_size change into mae.yaml – see the warning above.

• Longer training: raise n_epochs and patience, and run on SLURM through run_grid.py.

Beyond that: add your own architecture to neuraltrain as a BaseBrainModelConfig subclass (which also means adapting mae_module.py), or train in your own codebase and bring only the finished model back through the API below.

Evaluating a model of your own

mae.yaml works because the encoder lives in this repo. A model that lives in your own script needs no YAML here: evaluate_model() takes the built instance.

from neuralbench import check_model, evaluate_model

model = MyFoundationModel()   # built and pretrained however you like

print(check_model(model, "eeg", "motor_imagery"))
scores = evaluate_model(model, "eeg", "all", name="my-fm", debug=True)

Run check_model() before you queue anything: it pushes synthetic batches of the selection’s shapes through the model, so a shape bug surfaces in seconds rather than an hour into a real run. One instance serves every task in the selection, so the model must accept any channel count and window length, and read channel identity from a channel_positions argument to forward. It needs no classifier head.

See Evaluating your own model for suites, running on SLURM, and changing the adaptation protocol.

Next steps

• Track 1 – EEG-to-Image

• Track 2 – BCI decoding

• Track 3 – Sleep onset

• Track 4 – EMG-to-Pose

• How to Submit a Model