SleepMaMi

This is the code respository for SleepMaMi: A Universal Sleep Foundation Model for Integrating Macro- and Micro-structures (ICML 2026)

---

Prerequisites

A CUDA GPU is required for training and inference: the FlashAttention and Mamba (mamba-ssm) kernels used by the encoders have no CPU implementation. Use a GPU that FlashAttention 2 supports (Ampere or newer, e.g. A100 / RTX 30xx+).

conda create -n sleepmami python==3.10
conda activate sleepmami
pip install -r requirements.txt

# conda install -c "nvidia/label/cuda-12.1.1" cuda-nvcc cuda-cudart-dev libcublas-dev libcurand-dev # if necessary
pip install flash-attn==2.7.4.post1 --no-build-isolation --cache-dir /tmp/pip_cache
apt update
apt install libgl1-mesa-glx
pip install mamba-ssm==2.2.3.post2 --no-build-isolation

Data preprocessing

Please prepare pretraining and finetuning data, following README in data directory.

Pretraining

Micro-Encoder pretraining

After pretraining dataset is prepared, prepare a config file(*.yaml) in configs directory.

Example config files are provided in configs directory.

Use the following command to start pretraining:

python pretrain.py --config PT_6mod_9ds_CL

Macro-Encoder pretraining

Similar to Micro-Encoder pretraining, prepare a config file in configs directory.

Then use the following command to start Macro-Encoder pretraining

python pretrain_seq.py --config PTseq_6mod_9ds_CL_full_night

Downstream tasks

After pretraining is completed, finetuning on downstream tasks can be done with the following command:

python finetune.py --config FT_6mod_CL

The config file should be located in configs directory. ft_task variable should be set to specify the finetuning task (e.g. sleep_staging, apnea_seg, diag, etc.)

Pretrained checkpoints and Inference

inference.py loads a pretrained / finetuned checkpoint and exposes three entry points. Signal preprocessing reuses the same data/data_loader.py pipeline used during training, so inputs are identical to what the checkpoints were trained on.

Task	Backbone	Output	Ckpt
micro	ClassificationMultiConv1d (classification_multi)	per-epoch Micro feature	micro.safetensors
macro_micro	ClassificationMultiConv1dMamba (classification_multi_full_night)	full-night Macro + per-epoch Micro features	macro_micro.safetensors
sleep_staging	ClassificationMultiConv1dMamba + classification head (SHHS1)	predicted sleep stage per epoch	sleep_staging.safetensors

Input. Provide one .npy per modality holding the full night as a sequence of 30-s epochs, shape (N, sig_len) where sig_len = 3000 (3000 for the default 30 s @ 100 Hz). Pass a directory containing EEG.npy, EOG.npy, … via --npy_dir, or list paths explicitly with --npy.

# 1. Micro feature from the epoch-level Micro-Encoder
python inference.py --task micro \
    --checkpoint checkpoints/micro.safetensors \
    --npy_dir /path/to/night --num_epochs 20

# 2. Macro- and Micro-features from the full-night model
python inference.py --task macro_micro \
    --checkpoint checkpoints/macro_micro.safetensors \
    --npy_dir /path/to/night

# 3. Sleep staging (Macro + Micro features + classification head)
python inference.py --task sleep_staging \
    --checkpoint checkpoints/sleep_staging.safetensors \
    --npy_dir /path/to/night

# Explicit per-modality paths instead of a directory:
python inference.py --task sleep_staging \
    --npy EEG=EEG.npy EOG=EOG.npy ECG=ECG.npy RESP=RESP.npy EMG=EMG.npy SaO2=SaO2.npy

Notes:

• .npy inputs are normalized by default to match the training pipeline: SaO2 is scaled with (clip(x, 0.7, 1.0) - 0.7) / 0.3 (percent values are divided by 100 first), and other modalities are z-normalized per epoch (re-applied by RevIN internally). Pass --no_normalize if your .npy arrays are already preprocessed.
• Released checkpoints (checkpoints/*.safetensors) contain model weights only (no optimizer / scheduler state). --checkpoint defaults to checkpoints/<task>.safetensors; a Lightning .ckpt from your own finetune.py run can also be passed, or converted to the same format with train_utils.export_inference_weights.
• Inference must run on a CUDA GPU (--device cuda, the default when one is available); running on CPU fails in the FlashAttention kernel.
• --dtype bfloat16 reproduces the bf16-mixed precision used in training.
• micro / macro_micro only need an encoder (pretrained) checkpoint — the classification head is left uninitialized and unused. sleep_staging requires a finetuned checkpoint that contains the head.
• The functions in inference.py (extract_micro_features, extract_macro_micro_features, predict_sleep_stages, and the load_npy_epoch_inputs / load_npy_full_night_inputs helpers) can be imported and called directly from Python.

SDB segmentation changes (Updated at Oct-28 2026)

The ICML paper reported Macro-F1 for SleepMaMi, but F1 of sleep disoredred breathing class for baselines. To make them consistent, we updated the arXiv paper to report F1 of sleep disordered breathing class for all models. The repository is also udpated accordingly.

The updated SDB segmentation results are as following.

BibTeX

Please use the below bibtex when you reference our work!

@inproceedings{
park2026sleepmami,
title={SleepMaMi: A Universal Sleep Foundation Model for Integrating Macro- and Micro-structures},
author={Park, Keondo and Na, Younghoon and Choi, You Rim and Ryu, Hyunwoo and Shin, Hyun-Woo and Kim, Hyung-Sin},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=e9hvlHlUas}
}