Add SleepFM foundation and sleep-staging models - #1106

Summary

This PR adds support for SleepFM (https://github.com/zou-group/sleepfm-clinical), a multimodal polysomnography foundation model introduced in:

Thapa et al., “A multimodal sleep foundation model for disease prediction,” Nature Medicine (2026).
DOI: https://doi.org/10.1038/s41591-025-04133-4

The contribution introduces two public Braindecode models:

• SleepFM: the pretrained channel-agnostic PSG encoder with a trial-level prediction head.
• SleepFMStager: the released downstream architecture for patch-wise sleep-stage prediction.

Both models follow Braindecode’s EEGModuleMixin conventions and support the authors’ official base and sleep-staging checkpoints.

Links to #925.

Motivation

SleepFM was trained on multimodal PSG recordings and can process heterogeneous physiological channels, including EEG, EOG, respiratory signals, ECG, and EMG.

The upstream maintainers indicated that a contribution would be welcome and requested clarification about the baselines used in the paper. This PR therefore includes
both the model implementation and an explicit description of those baselines.

Model architecture

SleepFM

SleepFM processes raw signals with shape:

(batch, channels, time)

The reference configuration expects signals sampled at 128 Hz and uses non-overlapping 5-second patches, corresponding to 640 samples per patch.

The model contains:

1. A shared six-block convolutional tokenizer for every channel and patch.
2. Transformer-based attention pooling across available channels.
3. Fixed sinusoidal positional encoding.
4. A temporal Transformer encoder.
5. Temporal attention pooling.
6. A configurable trial-level prediction head.

The default parameters reproduce the released base configuration:

  SleepFM(
      patch_size=640,
      embed_dim=128,
      num_heads=8,
      num_layers=6,
      pooling_heads=8,
      drop_prob=0.3,
      max_seq_length=128,
  )

The model exposes:

  logits = model(x)
  result = model(x, return_features=True)
  tokens = model.tokenize(x)
  pooled, contextual_tokens = model.encode(x)

The returned shapes are:

  logits:             (batch, n_outputs)
  features:           (batch, embed_dim)
  tokens:             (batch, channels, patches, embed_dim)
  contextual_tokens:  (batch, patches, embed_dim)

SleepFMStager

SleepFMStager combines the released SleepFM tokenizer with the authors’ sleep-staging architecture:

1. Shared raw-signal tokenizer.
2. Attention pooling across channels.
3. Temporal Transformer encoder.
4. Bidirectional LSTM.
5. Patch-wise output projection.

Its output shape is:

(batch, n_outputs, n_patches)

With the official checkpoint, n_outputs=5 represents the five sleep-stage classes used upstream.

No softmax is included, allowing the logits to be consumed directly by the selected Braindecode or PyTorch loss.

Variable-channel masking

Both models accept an optional channel padding mask:

  output = model(x, channel_mask=channel_mask)

The mask has shape:

(batch, channels)

A value of True marks a missing or padded channel.

Masked channels are excluded from attention pooling and the subsequent masked mean. Each sample must contain at least one valid channel.

This allows examples with different PSG montages or modality availability to be batched together without treating padded channel values as real observations.

Input validation

The implementation validates:

• Input rank and expected (batch, channels, time) layout.
• Sampling frequency of exactly 128 Hz.
• Availability of at least one complete patch.
• Compatibility of patch_size with the six convolutional blocks.
• Divisibility of the embedding dimension by configured attention heads.
• Positional-encoding capacity.
• Channel-mask shape, type, and values.
• Samples with every channel masked.
• Compatibility of checkpoint keys and tensor shapes.

Trailing samples that do not form a complete patch are discarded, matching the upstream patching behavior.

Official checkpoint compatibility

The following loading methods are provided:

  model.load_pretrained_backbone(base_checkpoint)

and, for sleep staging:

  stager.load_pretrained_backbone(base_checkpoint)
  stager.load_pretrained_staging_head(staging_checkpoint)

The loaders support:

• Plain state dictionaries.
• Dictionaries nested under state_dict.
• Distributed module. prefixes.
• The upstream positional-encoding key format.
• Mapping the upstream staging fc layer to Braindecode’s final_layer.

Mappings are intentionally minimal. Unexpected or missing model keys are reported instead of silently discarding incompatible parameters.

Opt-in network tests download and load the official checkpoints:

  
  sleepfm/checkpoints/model_base/best.pt
  sleepfm/checkpoints/model_sleep_staging/best.pth

Both official-checkpoint tests pass.

Baselines used in the paper

The paper compares SleepFM with two supervised baselines.

Demographics baseline

The demographics-only baseline is a shallow MLP:

  age, sex, BMI, race/ethnicity
                │
             Linear(4, 32)
                │
               ReLU
                │
              Dropout
                │
         disease output layer

This baseline is not added as a separate Braindecode model because it does not process electrophysiological signals.

End-to-End PSG baseline

The End-to-End PSG baseline is trained jointly from random initialization and contains:

• The raw-signal convolutional tokenizer.
• Attention pooling across PSG channels.
• A bidirectional LSTM for temporal processing.
• An embedding of age and sex.
• A disease-prediction head.

This baseline does not use SleepFM self-supervised pretraining.

It is not equivalent to constructing SleepFM without pretrained weights: its downstream temporal architecture and demographic fusion are different. It is not included
as a separate public class in this PR because it requires nonelectrophysiological covariates and the paper’s disease-survival training pipeline.

The PR focuses on the reusable electrophysiological components and the released sleep-staging path.

Braindecode integration

The contribution includes:

• Public imports from braindecode.models.
• Entries in all.
• Model discovery and mandatory-parameter registration.
• Model summary metadata and parameter counts.
• Unified return_features behavior.
• reset_head(n_outputs) support.
• torch.export coverage.
• Reduced-configuration torch.compile/Dynamo graph-capture coverage.
• API autosummary entries.
• Foundation-model categorization.
• Release notes.
• License and attribution notice.

SleepFMStager is registered as a non-trial-classification model because it returns token-wise logits rather than a single prediction per recording.

Usage example

  import torch

  from braindecode.models import SleepFM

  model = SleepFM(
      n_chans=4,
      n_outputs=2,
      n_times=3840,
      sfreq=128,
  )

  x = torch.randn(8, 4, 3840)

  # True marks a missing or padded channel.
  channel_mask = torch.tensor(
      [
          [False, False, False, False],
          [False, False, True, True],
          [False, False, False, True],
          [False, False, False, False],
          [False, True, True, True],
          [False, False, False, True],
          [False, False, True, True],
          [False, False, False, False],
      ]
  )

  logits = model(x, channel_mask=channel_mask)
  features = model(
      x,
      channel_mask=channel_mask,
      return_features=True,
  )["features"]

  print(logits.shape)    # (8, 2)
  print(features.shape)  # (8, 128)

Sleep-staging example:

  import torch

  from braindecode.models import SleepFMStager

  model = SleepFMStager(
      n_chans=4,
      n_outputs=5,
      n_times=3840,
      sfreq=128,
  )

  x = torch.randn(8, 4, 3840)
  logits = model(x)

  # 3840 samples / 640 samples per patch = 6 patches
  print(logits.shape)  # (8, 5, 6)

Official checkpoint loading:

  import torch

  from braindecode.models import SleepFMStager

  base_checkpoint = torch.load("best.pt", map_location="cpu")
  staging_checkpoint = torch.load("best.pth", map_location="cpu")

  model = SleepFMStager(
      n_chans=4,
      n_outputs=5,
      n_times=3840,
      sfreq=128,
  )

  model.load_pretrained_backbone(base_checkpoint)
  model.load_pretrained_staging_head(staging_checkpoint)
  model.eval()

Testing

The implementation was tested with:

  pytest -q \
    test/unit_tests/models/test_sleepfm.py \
    test/unit_tests/models/test_return_features.py

Result:

68 passed, 3 skipped

Shared SleepFM integration coverage:

  pytest -q test/unit_tests/models/test_integration.py -k "SleepFM"

Result:

18 passed, 6 skipped

Official checkpoint tests:

  pytest -q \
    test/unit_tests/models/test_foundation_models.py \
    -k "sleepfm_official" \
    --run-network

Result:

2 passed

The standard repository quality gate was also run:

  pre-commit run -a

All configured hooks passed, including:

• Ruff formatting and linting.
• MyPy.
• isort.
• codespell.
• docstrfmt.
• Sphinx RST lint.
• File and whitespace checks.

git diff --check also passes.

Environmental verification notes

The complete model test directory was attempted, but the local Python 3.13/macOS environment encountered a native segmentation fault inside scikit-learn’s nearest-neighbor implementation. This occurred outside the SleepFM tests.

A complete Sphinx gallery build was also attempted. SleepFM autosummary pages were generated, but the full repository build was prevented by external intersphinx
network access, gallery dataset downloads, and sandbox multiprocessing restrictions. The configured RST, docstrfmt, and Sphinx-lint checks pass.

Licensing

The upstream SleepFM implementation and released checkpoints use the Creative Commons Attribution-NonCommercial 4.0 International license.

Consequently:

• braindecode/models/sleepfm.py retains upstream attribution.
• The file is explicitly identified as CC BY-NC 4.0.
• NOTICE.txt records the source repository, copyright attribution, and inherited noncommercial restriction.
• The implementation and official weights must not be treated as BSD-3-licensed components.

Upstream repository: https://github.com/zou-group/sleepfm-clinical
License: https://creativecommons.org/licenses/by-nc/4.0/

Files changed

• braindecode/models/sleepfm.py

  • SleepFM models, internal components, validation, and checkpoint mapping.

• braindecode/models/init.py

  • Public exports.

• braindecode/models/util.py

  • Model discovery and mandatory signal parameters.

• braindecode/models/summary.csv

  • Documentation metadata and parameter counts.

• test/unit_tests/models/test_sleepfm.py

  • Focused architecture and behavior tests.

• test/unit_tests/models/test_foundation_models.py

  • Official checkpoint tests.

• test/unit_tests/models/test_return_features.py

  • Unified feature and head-reset coverage.

• test/unit_tests/models/test_integration.py

  • Export, compile, and integration handling.

• docs/api.rst

  • Public API entries.

• docs/models/models_categorization.rst

  • Foundation-model categorization.

• docs/whats_new.rst

  • Release note and baseline explanation.

• NOTICE.txt

  • CC BY-NC attribution and licensing terms.

Checklist

• Added public model implementations.
• Followed EEGModuleMixin conventions.
• Added return_features and reset_head.
• Added validation and variable-channel masking.
• Added official checkpoint loaders.
• Tested official base and staging checkpoints.
• Added focused and shared integration tests.
• Added API and model-category documentation.
• Documented the paper’s baselines accurately.
• Added release notes.
• Added upstream attribution and licensing notice.
• Ran all configured pre-commit hooks.
• No new runtime dependency was introduced.