ProtoSleepNet

Prototype-based, interpretable sleep staging — where reading the prototypes is reading the decision.

Code for “Prototype-based interpretable sleep staging with physiologically plausible sub-stage pattern discovery” (ProtoSleepNet), under review at npj Digital Medicine. The revised manuscript is on the project website; version 1 is on Research Square, 10.21203/rs.3.rs-9169987/v1, under an earlier title.

▶ Try it live — protosleepnet-demo.pages.dev — an interactive atlas of the prototype space: click any epoch to see the exact evidence for its stage.

ProtoSleepNet wraps a sleep-staging backbone (SeqSleepNet or SleepTransformer) with a Prototype Sub-Stage (PSS) module: per-channel encoding, modality embeddings, accuracy-weighted modality dropout, a Transformer channel mixer and a dual-residual path, yielding contextualized epoch embeddings that are quantized into a small, interpretable prototype codebook. The prototypes carry physiologically plausible, AASM-coherent sub-stage patterns and support exploratory clinical probing (Alzheimer’s, Parkinson’s).

📖 Documentation: guidogagl.github.io/protosleepnet — method overview, a guide to how to interpret the model, the embedded live demo, and the full paper-reproduction recipe.

Built on physioex v2.0.2.

The model architecture, data loaders, preprocessing, dataset splits, seeds and evaluation metrics all live in physioex. This repository holds only the experiment pipeline that produced every figure and table in the paper and its Supplementary Notes.

Installation

conda env create -f environment.yml && conda activate protosleepnet   # optional
pip install -e .                     # installs the `protosleepnet` package + physioex==2.0.2
pip install -e ".[events]"           # + YASA/pyedflib, only for the detected-events analysis

Requires Python ≥ 3.10. After pip install -e . every script runs as a module from the repo root — one convention throughout:

python -m protosleepnet.train --help
python -m protosleepnet.figure_reconstruction.spectral_signature --help

Pretrained weights

Weights are on the HuggingFace Hub and load through physioex. Names follow the physioex convention <model>-<primary-author>. The paper’s models live in 4rooms/sleep-prototypes (the Phan originals {seqsleepnet,sleeptransformer}-phan are in 4rooms/physioex):

from physioex.models import load_from_pretrained
# PST — SleepTransformer backbone (SHHS, L=21)
model = load_from_pretrained("protosleeptransformer-gagliardi", repo_id="4rooms/sleep-prototypes")
# PSN — SeqSleepNet backbone (MASS, L=20)
# load_from_pretrained("protosleepnet-gagliardi", repo_id="4rooms/sleep-prototypes")
# input  (batch, L, 3, 29, 129) STFT log-power (EEG/EOG/EMG); output (batch, L, 5) AASM logits

The 3-channel baselines and ablation variants are published alongside them: {seqsleepnet,sleeptransformer}-gagliardi[-dropout|-mixer].

bash reproduce.sh runs a smoke path (load weights → forward pass → regenerate one figure).

Data

11 PSG datasets / 20 dataset-sources (12,317 subjects). Raw recordings are not redistributed here — obtain them from the sources below and preprocess them with physioex (this repo depends on physioex’s preprocessing; it ships none of its own). See the physioex docs and the launchers in examples/slurm/.

Datasets	Access
SHHS, MrOS, MESA, WSC, HomePAP	NSRR — controlled/registered (https://sleepdata.org)
HMC, Sleep-EDF	PhysioNet — open
MASS, DCSM	open on request
ASD (Alzheimer), KPD (Parkinson)	KU Leuven — restricted (contact authors; ethics S61792/S70708)

Configuration (environment variables)

Scripts read paths from the environment so no absolute paths are baked in:

Variable	Meaning	Default
PHYSIOEX_DATA	physioex dataset root (preprocessed datasets; raw MASS SS2 EDFs for the detected-events analysis)	—
PROTOSLEEPNET_DATA / EXPERIMENT_DIR	root for predictions, embeddings and reconstruction outputs	data
PROTOSLEEPNET_MODELS	checkpoints / training statistics	models
PROTOSLEEPNET_RECON_M12	per-prototype reconstruction records (Fig 3–4, Tab 1–2 sources)	data/reconstructions/M12
PROTOSLEEPNET_OUTPUTS	root of the analysis modules (epoch-embeddings/, json/, models/, analysis/)	outputs/protosleepnet
PROTOSLEEPNET_KPD_RAW	raw KPD recordings (lights-off markers; clinical probing only)	—

Cluster launchers additionally source examples/slurm/env.sh (copy from examples/slurm/env.sh.example).

Repository layout

src/protosleepnet/            # installable package (run as python -m protosleepnet.<mod>)
  train.py  test_pretrained.py  extract_embeddings.py   # staging training / evaluation
  learn_prototypes_vq.py  test_prototypes.py            # VQ codebook learning + quantized inference
  test_occlusion.py  test_prototype_randomization.py    # occlusion robustness, prototype randomization
  seed_*.py  convert_and_predict_seeds.py              # seed-stability pipeline
  build_protosleepnet.py                                # shared model factory
  ablation/  baselines/                                 # component ablation; residual / mixer / dropout variants
  posthoc_prototypes/                                   # k-means codebooks, quantization disagreement, codebook refinement
  proto_reconstruction/                                 # data-driven / model-driven / hybrid reconstruction
  figure_reconstruction/                                # rules, spectral & relevance signatures, local IG,
                                                        #   threshold sweeps, stage permutation, detected events
  clinical_probing/                                     # atlas, per-subject records, 20 features, nested LOSO,
                                                        #   hypnogram comparator, DeLong inference, deficit structure
  plot/                                                 # portable figure scripts (from predictions)
  demo/                                                 # interactive-demo precompute pipeline + static app
examples/slurm/    # cluster launchers; copy env.sh.example -> env.sh (analysis/ = CPU-only analyses)
data/              # small committed figure-source JSON only (see below)
scripts/           # HuggingFace upload of the released weights + model cards
tests/             # smoke tests (import graph + optional pretrained forward pass)

Large artifacts (predictions, embeddings, reconstruction arrays, checkpoints) are not in git. Only small figure-source JSON is committed under data/: data/reconstructions/M12/ (per-prototype records behind Fig 3–4 and Tab 1–2, also the input of the demo and of the checklist analyses) and data/probing_staging_summary/ (staging tables of Supplementary Note 1).

Reproducing the paper (Figure/Table → experiment → command)

Each command below is the entry point; pass --help for its options. Full experiment settings (M value, seeds, datasets) are documented per script. Modules under clinical_probing/ run in the order listed.

Paper asset	Experiment	Command
Fig 2a; Supp Note 1	In-domain + OOD staging, per-class, confusion, stats	python -m protosleepnet.train · python -m protosleepnet.test_pretrained · python -m protosleepnet.extract_embeddings
Supp Note 2	Component ablation (4 variants) + occlusion robustness	python -m protosleepnet.ablation.train_ablation --variant {baseline,dropout,mixer,protosleepnet} (SleepTransformer) / ...train_ablation_seqsleepnet (SeqSleepNet) · python -m protosleepnet.test_occlusion · python -m protosleepnet.plot.occlusion · python -m protosleepnet.plot.occlusion_ablation
Fig 2b; Supp Note 3	M-sweep, VQ methods, residual baselines	python -m protosleepnet.learn_prototypes_vq · python -m protosleepnet.test_prototypes · python -m protosleepnet.posthoc_prototypes.{extract_epoch_embeddings,learn_prototypes_vq_kmeans,test_posthoc_prototypes,test_vq_paired,probe_posthoc_prototypes} · python -m protosleepnet.baselines.{train_*_residual,extract_residual_embeddings,learn_residual_prototypes_vq,test_residual_prototypes} · python -m protosleepnet.plot.{vq,residual,residual_ablation}
Supp Note 3 (§5)	Continuous-vs-quantized disagreement; paired per-subject κ	python -m protosleepnet.posthoc_prototypes.quantization_disagreement · python -m protosleepnet.posthoc_prototypes.non_inferiority
Fig 4; Supp Notes 5, 6	Prototype reconstruction (data/model/hybrid) + cross-dataset	python -m protosleepnet.proto_reconstruction.{data_driven,model_driven,hybrid} · python -m protosleepnet.figure_reconstruction.compute_cross_dataset_metrics
Fig 3; Tab 1, 2; Supp Notes 4, 6	Codebook summaries, band-ablation rules, coherence checklist, local IG	python -m protosleepnet.figure_reconstruction.{rule_learning,spectral_signature,relevance_signature,combinatorial_ablation,compute_local_explanations}
Tab 2; Supp Note 6 (§5)	Checklist threshold sweeps + joint margins; stage-permutation control	python -m protosleepnet.figure_reconstruction.threshold_sweep · python -m protosleepnet.figure_reconstruction.stage_permutation
Supp Note 6 (§10)	Automatically detected spindles / slow waves vs band relevance (MASS SS2; needs [events])	python -m protosleepnet.figure_reconstruction.detect_events · python -m protosleepnet.figure_reconstruction.events_relevance
Supp Note 8	Prototype randomization	python -m protosleepnet.test_prototype_randomization
Supp Note 9	Seed stability (training seeds, reconstruction, MMD)	python -m protosleepnet.seed_stability_pipeline · python -m protosleepnet.convert_and_predict_seeds · python -m protosleepnet.seed_recon_all_m · python -m protosleepnet.seed_mmd_analysis · python -m protosleepnet.seed_stability_summary
Supp Note 9 (§4)	Coarse-to-fine codebook refinement across adjacent M	python -m protosleepnet.posthoc_prototypes.codebook_refinement
Fig 5; Tab 3; Supp Note 7	Clinical probing (ASD, KPD): frozen embeddings → training atlas → per-subject records → 20 features → 34-descriptor comparator → nested LOSO → DeLong inference → deficit structure	python -m protosleepnet.clinical_probing.staging.extract_embeddings · python -m protosleepnet.clinical_probing.atlas · python -m protosleepnet.clinical_probing.subject_stats · python -m protosleepnet.clinical_probing.build_features · python -m protosleepnet.clinical_probing.hypnogram_comparator · python -m protosleepnet.clinical_probing.nested_loso · python -m protosleepnet.clinical_probing.inference · python -m protosleepnet.clinical_probing.structure
Live demo	Subject selection, anonymized subset, PaCMAP atlas, per-epoch IG, prototype cards	python -m protosleepnet.demo.{select_subjects,precompute,precompute_indomain,build_cards,subset_bundle,evaluate_explanations}

CPU-only analyses have ready-made launchers under examples/slurm/analysis/. Manuscript figure/table assembly (LaTeX-coupled emitters) is kept out of this repo, with the paper sources.

Tests

pip install -e . pytest
pytest                 # import-graph smoke test (no network)
pytest --runslow       # also pull weights from HuggingFace and run a forward pass

Citation

If you use this code, please cite the paper and the software — see CITATION.cff. This work builds on physioex v2.0.2. License: MIT (LICENSE).