OSF: On Pre-training and Scaling of Sleep Foundation Models

🔥 News

• [2026-2-24] Our codebase and checkpoint is released.
• [2026-2-22] Our paper is out.

📖 Introduction

Polysomnography (PSG) provides the gold standard for sleep assessment but suffers from substantial heterogeneity across recording devices and cohorts. There have been growing efforts to build general-purpose foundation models (FMs) for sleep physiology, but lack an in-depth understanding of the pre-training process and scaling patterns that lead to more generalizable sleep FMs. To fill this gap, we curate a massive corpus of 166,500 hours of sleep recordings from nine public sources and establish SleepBench, a comprehensive, fully open-source benchmark. Leveraging SleepBench, we systematically evaluate four families of self-supervised pre-training objectives and uncover three critical findings: (1) existing FMs fail to generalize to missing channels at inference; (2) channel-invariant feature learning is essential for pre-training; and (3) scaling sample size, model capacity, and multi-source data mixture consistently improves downstream performance. With an enhanced pre-training and scaling recipe, we introduce OSF, a family of sleep FMs that achieves state-of-the-art performance across nine datasets on diverse sleep and disease prediction tasks. Further analysis of OSF also reveals intriguing properties in sample efficiency, hierarchical aggregation, and cross-dataset scaling.

📖 Table of Contents

1. Installation
2. Quick Start
3. Pretrained Weights
4. Usage
5. Benchmark Evaluations
6. Supported Datasets
7. Citation

💿 Installation

git clone https://huggingface.co/yang-ai-lab/OSF-Base
cd OSF-Base
conda env create -f environment.yml
conda activate myenv

Dependencies

• Python >= 3.10
• PyTorch >= 2.9.0
• PyTorch Lightning >= 2.5.5

🚀 Quick Start

We provide a demo notebook (demo.ipynb) demonstrating how to extract embeddings from PSG signals using the pretrained model.

import torch
from osf.backbone.vit1d_cls import vit_base

# Load pretrained weights (included in this repo)
payload = torch.load("osf_backbone.pth", map_location="cpu", weights_only=False)
meta = payload["metadata"]

# Initialize model
backbone = vit_base(
    num_leads=meta["num_leads"],        # 12 channels
    seq_len=meta["seq_len"],            # 1920 (64 Hz × 30 s)
    patch_size=meta["patch_size_time"],
    lead_wise=meta["lead_wise"],
    patch_size_ch=meta["patch_size_ch"],
)
backbone.load_state_dict(payload["state_dict"])
backbone.eval()

# Extract embeddings
# x: [B, 12, 1920] - 12-channel PSG, 64 Hz × 30 seconds
with torch.no_grad():
    cls_embs, patch_embs = backbone.forward_encoding(x, return_sequence=False)
# cls_embs: [B, 768] - Global epoch-level representation
# patch_embs: [B, 90, 768] - Local patch representations

📦 Pretrained Weights

The pretrained weights are included in this repository. You can download them via the Hugging Face Hub:

from huggingface_hub import hf_hub_download
checkpoint_path = hf_hub_download(repo_id="yang-ai-lab/OSF-Base", filename="osf_backbone.pth")

Or via the CLI:

huggingface-cli download yang-ai-lab/OSF-Base osf_backbone.pth

👩‍💻 Usage

Input Format

Expected input format:

• 12 PSG Channels: ECG, EMG_Chin, EMG_LLeg, EMG_RLeg, ABD, THX, NP, SN, EOG_E1_A2, EOG_E2_A1, EEG_C3_A2, EEG_C4_A1
• Sample Rate: 64 Hz
• Epoch Length: 30 seconds
• Input Shape: [B, 12, 1920]

Pretraining

We support multiple self-supervised pretraining methods, for example, to launch pre-training of our OSF method, run pretraining:

python main_pretrain.py \
    --model_name "dino_ours" \
    --psg_encoder_name "vit_base" \
    --batch_size 256 \
    --lr 5e-5 \
    --max_epochs 30 \
    --num_devices 4 \
    --patch_size_time 64 \
    --patch_size_ch 4 \
    --precision "bf16-mixed"

See main_pipleines/main_pretrain.py for more detailed settings.

Fine-tuning

Fine-tune the pretrained model on downstream tasks:

python main_finetune.py \
    --model_name "dino_ours" \
    --ckpt_path "/path/to/pretrained/checkpoint.ckpt" \
    --downstream_dataset_name "shhs" \
    --eval_label "Stage" \
    --train_data_pct 1.0 \
    --max_steps 500 \
    --lr 0.1 \
    --num_devices 4

📊 Benchmark Evaluations

Benchmarked SSL Methods

Downstream Tasks

Epoch-level Classification Tasks:

Evaluation Settings

For example scripts, see main_pipelines and bash_scripts folders.

📊 SleepBench

We aggregated nine large-scale datasets from the National Sleep Research Resource platform to form our SleepBench. To acquire these resources, please accesss the following links. We will open-source our preprocessing codes in later version.

For new users, please apply for an account and access to each of these datasets following instructions here NSRR Registration

📁 Project Structure

OSF-Open-Sleep-Foundation-Model/
├── osf/
│   ├── backbone/          # ViT backbone implementations
│   │   └── vit1d_cls.py
│   ├── models/            # SSL model implementations
│   │   └── dino_model_cls.py
│   │   
│   ├── datasets/          # Data loading utilities
│   └── utils/             # Helper functions
├── main_pipelines/        # Training scripts
│   ├── main_pretrain.py
│   └── ...
├── bash_scripts/          # Example bash scripts
├── osf_backbone.pth       # Pretrained model weights
├── demo.ipynb             # Quick start demo
├── config.py              # Dataset and channel configurations
└── train_config.py        # Training configurations

📝 Citation

If you use this code or models in your research, please cite our paper:

@article{shuai2026osf,
  title={OSF: On Pre-training and Scaling of Sleep Foundation Models},
  author={Shuai, Zitao and Xu, Zongzhe and Yang, David and Wang, Wei and Yang, Yuzhe},
  journal={arXiv preprint arXiv:2603.00190},
  year={2026}
}