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

This paper introduces OSF, a family of state-of-the-art sleep foundation models trained on a massive 166,500-hour corpus and evaluated via the open-source SleepBench, which demonstrates that channel-invariant pre-training and systematic scaling of data, capacity, and diversity significantly enhance generalization across diverse sleep and disease prediction tasks.

The Gist

Imagine you want to teach a computer to understand human sleep. To do this, you need to show it millions of nights of sleep recordings. These recordings, called Polysomnography (PSG), are like a symphony of biological signals: brain waves, heartbeats, breathing patterns, and muscle movements.

However, there's a big problem. Just like every orchestra has different instruments, every sleep study uses different equipment. Some hospitals record 12 channels of data; some home devices only record 2. Sometimes sensors fall off in the middle of the night. If you train a computer to only understand a full orchestra, it will be confused when it hears just a solo violin.

This paper introduces OSF (Open Sleep Foundation Model), a new, super-smart AI designed to understand sleep no matter what "instruments" are playing. Here is the story of how they built it, explained simply.

1. The Problem: The "Missing Instrument" Crisis

The researchers found that previous sleep AI models were like musicians who could only play if every single instrument in the orchestra was present. If you took away the brain sensors (the violins), the model couldn't tell if someone was asleep or awake. If you took away the breathing sensors (the drums), it couldn't detect sleep apnea.

In the real world, data is messy. Sensors break, and home devices are simpler. The old models failed when faced with these "missing instruments."

2. The Solution: Building a Massive Library (SleepBench)

To fix this, the team didn't just look at one dataset. They built SleepBench, a massive digital library containing 166,500 hours of sleep recordings from nine different public sources.

Think of this as gathering recordings from 21,000 different people, using 9 different types of recording equipment, covering all kinds of ages and health conditions. This gave them a diverse "training camp" to teach the AI to be flexible.

3. The Three Big Discoveries

By testing different ways to train the AI, they found three golden rules for building a sleep expert:

• Rule #1: Don't rely on a single instrument.
  They discovered that if you train the AI to expect all sensors to be there, it fails when they aren't. The AI needs to learn that the meaning of sleep is the same, whether it's hearing a full orchestra or just a few instruments.
• Rule #2: Teach the AI to be "Channel-Invariant."
  This is the secret sauce. They taught the AI by randomly "muting" different sensors during training. Imagine playing a song and telling the AI, "Okay, pretend the drums are gone. What does the music sound like now?" or "Pretend the violins are gone."
  By forcing the AI to learn the essence of sleep regardless of which sensors are active, it became incredibly robust. It learned that a sleeping brain looks a certain way, even if the breathing sensor is missing.
• Rule #3: Bigger is better (if you do it right).
  They found that simply feeding the AI more data and giving it a bigger brain (more computing power) made it smarter. But this only worked because of Rule #2. If you just give a rigid AI more data, it hits a wall. If you give a flexible AI (trained with the "mute" technique) more data, it keeps getting smarter and smarter.

4. The Result: OSF

The result is OSF, a new "Foundation Model" for sleep. Think of OSF as a master sleep detective who has studied every type of recording device imaginable.

• It's a chameleon: It works perfectly on hospital-grade machines with 12 sensors, and it still works well on simple home headbands with only 2 sensors.
• It's a doctor: It doesn't just tell you if you are asleep; it can predict if you have heart disease, diabetes, or high blood pressure based on your sleep patterns.
• It's efficient: It learns faster and needs fewer examples to get good at a new task compared to older models.

The Analogy: The Universal Translator

Imagine you are trying to learn a language.

• Old Models were like students who only learned to speak English if they had a dictionary, a grammar book, and a native speaker standing right next to them. If you took away the dictionary, they froze.
• OSF is like a polyglot who learned the language by listening to radio, watching TV, talking to strangers, and reading street signs. If you take away the radio, they can still understand the street signs. They learned the concept of the language, not just the specific tools used to teach it.

Why This Matters

This is a huge step forward for health. It means we can finally use simple, cheap, home-based sleep trackers to diagnose serious diseases with the same accuracy as expensive hospital tests. It makes high-quality sleep medicine accessible to everyone, not just those with access to a full hospital lab.

In short: The researchers taught a computer to understand sleep by making it practice with "broken" data, resulting in a model that is tougher, smarter, and ready for the real world.

Technical Summary

1. Problem Statement

Polysomnography (PSG) is the gold standard for sleep assessment but faces significant challenges in real-world deployment:

• Heterogeneity: Recordings vary widely across different devices, patient cohorts, and acquisition protocols.
• Channel Incompleteness: In practice, specific channels (e.g., EEG for brain activity or respiratory sensors) are often missing due to home-based studies, sensor dropouts, or specific study designs.
• Lack of Generalization: Existing Sleep Foundation Models (FMs) often fail to generalize when tested on data with missing channels or different demographic distributions.
• Unknown Scaling Laws: There is a lack of systematic understanding regarding how pre-training objectives, data volume, model capacity, and multi-source data mixing affect the performance of sleep FMs.

2. Methodology

A. SleepBench: A Comprehensive Benchmark

To address the lack of standardized evaluation, the authors curated SleepBench, the largest fully open-source sleep benchmark to date.

• Scale: Aggregates 166,500 hours of PSG recordings from 21,482 sleep studies across 9 public datasets (SHHS, NCHSDB, CFS, CCSHS, WSC, MROS, MESA, CHAT, SOF).
• Standardization: Utilizes a standardized 12-channel montage covering four physiological modalities:
  • Brain: EEG (C3-A2, C4-A1) and EOG (E1-A2, E2-A1).
  • Respiration: Abdominal, Thorax, Nasal Pressure, Snore.
  • Cardiac: ECG.
  • Somatic: Chin EMG, Left Leg EMG, Right Leg EMG.
• Preprocessing: Includes manual quality control, per-night z-score normalization, and segmentation into non-overlapping 30-second epochs (yielding ~20 million epochs).

B. Systematic Pre-training Analysis

The authors conducted a controlled study evaluating four families of self-supervised learning (SSL) objectives:

1. Contrastive Learning: SimCLR.
2. Reconstruction-based: MAE, VQ-VAE.
3. Autoregressive Modeling: Autoregression.
4. Self-Distillation: DINO.

They investigated three critical design dimensions:

1. Missing Channel Robustness: Simulating real-world scenarios where specific channels (e.g., brain or respiratory) are zero-masked at inference time.
2. Augmentation Strategies: Comparing standard time-wise masking against channel masking (randomly dropping 50% of channels during pre-training).
3. Scaling Laws: Analyzing the impact of increasing pre-training sample size, model capacity (ViT-1M to ViT-85M), and multi-source data mixing.

C. OSF (Open Sleep Foundation Model)

Based on the findings, the authors introduced OSF, a family of sleep FMs.

• Architecture: Based on a ViT-85M Transformer encoder.
• Core Recipe: Uses DINO (self-distillation) as the base objective but incorporates a specific two-stage masking augmentation:
  1. Channel Masking: Randomly drops 50% of input channels to force channel-invariant feature learning.
  2. Temporal Block Masking: Masks contiguous time segments (30-60% ratio).
• Training: Pre-trained on the full multi-source SleepBench corpus with explicit scaling of model size and data diversity.

3. Key Findings

The paper establishes three critical findings that guide the construction of robust sleep FMs:

1. Existing FMs Fail Under Missing Channels: Current state-of-the-art models suffer significant performance drops when tested on data with missing channels (e.g., no EEG for sleep staging). They are not robust to input incompleteness.
2. Channel-Invariant Learning is Essential: Explicitly encouraging the model to learn features that are invariant to the available channel subset (via channel masking during pre-training) is crucial. This significantly improves robustness and transferability, particularly for contrastive and distillation-based methods.
3. Scaling Laws Hold for Sleep: Unlike previous observations where sleep models saturated quickly, OSF demonstrates consistent performance gains when scaling:
   • Data: Increasing pre-training sample size improves downstream performance.
   • Model: Larger model capacities (up to 85M parameters) yield better representations without overfitting.
   • Diversity: Multi-source data mixing outperforms single-source pre-training, enhancing generalization across cohorts.

4. Results

OSF was evaluated across 9 datasets and 8 downstream tasks (4 epoch-level sleep tasks and 3 patient-level disease predictions).

• State-of-the-Art Performance: OSF consistently outperforms all baselines (including SleepFM, SimCLR, DINO, MAE, VQ-VAE) in both Linear Probing and Full Fine-tuning settings.
  • Sleep Staging: Achieves 97.3% AUC (Linear Probing) and 97.9% AUC (Fine-tuning) on MROS, surpassing the previous best (SleepFM) by a clear margin.
  • Event Detection: Sets new records for Arousal, Hypopnea, Oxygen Desaturation, and Central Apnea detection.
• Robustness to Missing Channels:
  • In "Head-band only" scenarios (Brain signals only), OSF improves sleep staging AUC by +0.6% over SleepFM.
  • In "In-home study" scenarios (Respiratory + ECG only), OSF improves Hypopnea detection AUC by +3.6%.
• Few-Shot Adaptation: OSF demonstrates superior sample efficiency, achieving the best performance in 1-shot, 5-shot, and 50-shot settings across all tasks.
• Disease Prediction: OSF shows improved transferability for patient-level disease classification (Coronary Disease, Diabetes, Hypertension), achieving the highest AUC and AUPRC scores.
• Embedding Quality: UMAP visualizations reveal that OSF learns more separable and structured clusters aligned with sleep stages compared to previous models.

5. Significance and Contributions

• First Systematic Study of Sleep FM Scaling: This work moves beyond reporting single-model gains to provide a rigorous analysis of why certain pre-training choices work, establishing scaling laws specific to sleep physiology.
• Solving the Missing Channel Problem: By identifying channel-invariant learning as a prerequisite for robustness, the paper provides a practical solution for deploying sleep models in diverse, real-world clinical and home settings where sensor availability is inconsistent.
• Open Science: The release of SleepBench (166k hours of data) and the OSF code/checkpoints democratizes access to high-quality sleep data and models, setting a new standard for reproducibility in the field.
• Practical Recipe: The paper offers a concrete "recipe" for building sleep FMs: use large-scale multi-source data, scale model capacity, and explicitly employ channel masking during pre-training.

In conclusion, OSF represents a significant leap forward in sleep foundation models, transforming them from fragile, dataset-specific tools into robust, generalizable systems capable of handling the heterogeneity of real-world sleep medicine.