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

Este artículo presenta OSF, una familia de modelos fundacionales para el sueño que, tras analizar 166.500 horas de registros y establecer el benchmark SleepBench, demuestra que el aprendizaje de características invariante a los canales y la escalabilidad de los datos mejoran significativamente la generalización y el rendimiento en tareas de predicción de sueño y enfermedades.

Lo esencial

¡Claro que sí! Imagina que el sueño es como una orquesta gigante donde cada instrumento (el cerebro, el corazón, la respiración, los músculos) toca su propia parte para crear la música de la noche.

Hasta ahora, los científicos tenían un problema: para entender esta música, necesitaban escuchar todos los instrumentos a la vez. Pero en la vida real, a veces los sensores se caen, o el paciente está en casa y solo tiene un dispositivo simple que no graba el cerebro. Es como intentar adivinar la canción de una orquesta si solo te dan la partitura de los tambores y te quitan los violines.

Aquí es donde entra este nuevo trabajo, llamado OSF (Modelo de Fundación del Sueño Abierto). Los autores (un equipo de UCLA y Emory) han creado algo revolucionario. Vamos a desglosarlo con analogías sencillas:

1. El Gran Archivo de Música (SleepBench)

Antes, los científicos tenían pequeños archivos de música de diferentes orquestas, pero no podían compararlos bien porque cada uno grababa cosas distintas.

• La solución: Crearon SleepBench, una "biblioteca de música" masiva y abierta. Recopilaron 166,500 horas de grabaciones de sueño de 9 fuentes diferentes. Es como tener un archivo de audio con millones de noches de sueño de todo tipo de personas.
• El objetivo: Usar esta biblioteca para entrenar a una "IA maestra" que entienda el sueño mejor que cualquier humano.

2. El Problema de los Instrumentos que Faltan

Los modelos anteriores (como SleepFM) eran como músicos que solo sabían tocar si tenían todos los instrumentos conectados. Si en la prueba (inferencia) faltaba un cable (por ejemplo, el del cerebro), el modelo se confundía y fallaba estrepitosamente.

• El hallazgo: Descubrieron que los modelos actuales no son lo suficientemente "flexibles". Si les quitas la señal del cerebro, no pueden adivinar la fase del sueño.

3. La Receta Secreta: "Aprender sin depender de un solo instrumento"

Para arreglar esto, probaron cuatro formas diferentes de enseñar a la IA. La clave fue una técnica llamada "enmascaramiento de canales".

• La analogía: Imagina que estás entrenando a un detective. En lugar de darle todas las pistas siempre, a veces le quitas la mitad de las pistas (ciertos sensores) y le obligas a resolver el caso con lo que queda.
• El resultado: Al entrenar a la IA así, aprendió a ser invariante al canal. Es decir, aprendió que "el corazón late rápido cuando hay estrés" o "la respiración se detiene cuando hay apnea", sin importar si también tiene la señal del cerebro o no. ¡Se volvió un detective infalible incluso con pistas incompletas!

4. La Ley de la Escala: Más es Mejor (pero con la receta correcta)

Otro descubrimiento importante fue sobre el tamaño.

• El mito: Pensaban que si entrenabas a un modelo con más datos, siempre mejoraría. Pero los modelos viejos se "cansaban" (se estancaban) después de cierto punto.
• La verdad: Con la nueva receta (entrenar con datos de muchas fuentes mezcladas y usar la técnica de "falta de pistas"), la IA sí mejora cuanto más la alimentas.
• La analogía: Es como estudiar para un examen. Si estudias con un solo libro viejo, llegas a un límite. Pero si estudias con una biblioteca entera de libros de diferentes autores y te obligas a practicar sin mirar las respuestas, tu conocimiento crece sin parar.

5. OSF: El Nuevo Campeón

Con todo esto, crearon OSF.

• ¿Qué hace? Es un modelo que puede predecir si alguien tiene apnea, diabetes o hipertensión solo mirando su sueño, y lo hace mejor que nadie.
• ¿Por qué es especial?
  • Es robusto: Si un sensor falla en casa, OSF sigue funcionando bien.
  • Es eficiente: Aprende mucho con pocos ejemplos (ideal para enfermedades raras).
  • Es general: Funciona bien en diferentes poblaciones y tipos de equipos.

En resumen

Este papel nos dice que para crear una Inteligencia Artificial que entienda el sueño humano de verdad, no basta con tener muchos datos. Necesitas:

1. Muchos datos (de muchas fuentes).
2. Un modelo grande (inteligente).
3. Entrenarlo de forma inteligente: Obligarlo a aprender cuando le faltan piezas del rompecabezas, para que cuando llegue al mundo real (donde los sensores siempre fallan un poco), no se rinda.

¡Es como pasar de tener un mapa incompleto a tener una brújula que funciona incluso si se te cae la mitad del mapa!

Resumen técnico

Resumen Técnico: OSF (Open Sleep Foundation Models)

1. Problema y Motivación

La polisomnografía (PSG) es el estándar de oro para la evaluación del sueño, pero presenta desafíos significativos para el desarrollo de modelos de aprendizaje automático generalizables:

• Heterogeneidad: Los registros varían enormemente entre dispositivos, cohortes de pacientes y protocolos de estudio.
• Incompletitud de Canales: En la práctica clínica y en estudios domiciliarios, a menudo faltan canales críticos (por ejemplo, señales cerebrales EEG/EOG en dispositivos portátiles o señales respiratorias en estudios de micro-arquitectura). Los modelos existentes fallan al inferir con canales faltantes.
• Falta de Escalabilidad: No existe una comprensión profunda de cómo las decisiones de pre-entrenamiento (objetivos de aprendizaje, aumento de datos, mezcla de fuentes) afectan la capacidad de generalización de los Modelos Fundacionales (FMs) para el sueño.
• Benchmarks Cerrados: La falta de un benchmark abierto, estandarizado y masivo limita la comparación justa entre diferentes métodos de aprendizaje auto-supervisado (SSL).

2. Metodología

A. SleepBench: Un Nuevo Benchmark Abierto
Los autores curaron un corpus masivo de 166,500 horas de registros de sueño procedentes de 9 fuentes públicas (incluyendo SHHS, MESA, CHAT, etc.), abarcando más de 21,000 estudios nocturnos.

• Estandarización: Se utiliza un montaje estandarizado de 12 canales agrupados por modalidad fisiológica: Cerebro (EEG/EOG), Respiración, Cardíaco (ECG) y Somático (EMG).
• Preprocesamiento: Se aplica control de calidad manual, normalización z-score por noche (con ajustes específicos para canales respiratorios) y segmentación en épocas de 30 segundos (aprox. 20 millones de épocas).
• División: Los datos se dividen en cohortes "in-domain" (para pre-entrenamiento y ajuste fino) y "out-of-domain" (OOD) para evaluación estricta de generalización.

B. Estudio Sistemático de Pre-entrenamiento
Se evaluaron cuatro familias de objetivos de aprendizaje auto-supervisado:

1. Aprendizaje Contrastivo (SimCLR).
2. Métodos basados en Reconstrucción (MAE, VQ-VAE).
3. Modelado Autoregresivo.
4. Auto-distilación (DINO).

C. Hallazgos Críticos y Diseño de OSF
A través de experimentos controlados, los autores identificaron tres principios clave que guían el diseño de OSF (Open Sleep Foundation Model):

1. Fallo ante canales faltantes: Los FMs existentes no generalizan bien cuando se eliminan canales durante la inferencia.
2. Invarianza de Canales: Es crucial entrenar al modelo para aprender características invariantes al conjunto de canales disponibles. Se implementa una estrategia de enmascaramiento de canales (dropping aleatorio del 50% de los canales) junto con el enmascaramiento temporal durante el pre-entrenamiento.
3. Leyes de Escalado: El rendimiento mejora consistentemente al aumentar el tamaño de la muestra, la capacidad del modelo y la mezcla de datos de múltiples fuentes.

D. Arquitectura de OSF

• Base: Utiliza un codificador Transformer (ViT) con 85 millones de parámetros.
• Objetivo: Se basa en DINO (auto-distilación) con las recetas de pre-entrenamiento mejoradas (enmascaramiento de canales + temporal).
• Estrategia de Agregación: Para tareas a nivel de paciente (enfermedades), se propone una selección de los top-k epochs relevantes en lugar del promedio simple, mejorando la agregación jerárquica.

3. Resultados Clave

A. Rendimiento en Tareas de Análisis del Sueño (Nivel de Época)
OSF alcanza el estado del arte (SOTA) en todas las tareas evaluadas (clasificación de etapas de sueño, detección de arousal, hipopnea, desaturación de oxígeno y apnea central) en 9 conjuntos de datos diferentes.

• Comparación: Supera consistentemente a modelos anteriores como SleepFM, SimCLR y DINO estándar, tanto en Linear Probing como en Fine-tuning completo.
• Ejemplo: En la tarea de clasificación de etapas de sueño (Sleep Staging) en el conjunto MROS (OOD), OSF logra un AUC de 97.3 frente a 96.4 de SleepFM.

B. Robustez ante Canales Faltantes (Inferencia)

• OSF demuestra una superioridad notable cuando se eliminan canales críticos (ej. solo señales respiratorias o solo cerebrales).
• Mientras que los modelos base sufren caídas drásticas de rendimiento, OSF mantiene una alta precisión al aprender representaciones invariantes, demostrando una mayor adaptabilidad a dispositivos reales con sensores limitados.

C. Predicción de Enfermedades (Nivel de Paciente)
En tareas de predicción de enfermedades cardiovasculares, diabetes e hipertensión, OSF supera a los métodos existentes, logrando mejoras significativas en AUC y AUPRC. Esto indica que las representaciones aprendidas capturan mejor la fisiología subyacente relacionada con patologías crónicas.

D. Eficiencia de Muestra y Escalado

• Few-Shot: OSF muestra una capacidad de adaptación superior en escenarios con pocos ejemplos (1-shot, 5-shot).
• Escalado: Se confirma que el rendimiento escala positivamente con el tamaño del modelo (hasta 85M parámetros) y la cantidad de datos de pre-entrenamiento, especialmente cuando se utiliza una mezcla de múltiples fuentes de datos.

4. Contribuciones Principales

1. SleepBench: La creación del benchmark de sueño de código abierto más grande y diverso hasta la fecha, facilitando la investigación reproducible.
2. Análisis Sistemático: El primer estudio exhaustivo que desentraña qué estrategias de pre-entrenamiento (objetivos SSL, aumentos de datos, escalado) realmente impulsan la generalización en modelos de sueño.
3. Descubrimiento de Invarianza de Canales: Demostración empírica de que el enmascaramiento de canales durante el pre-entrenamiento es esencial para la robustez en inferencia con datos incompletos.
4. OSF (Open Sleep Foundation Model): La liberación de una familia de modelos pre-entrenados que establece nuevos estándares de rendimiento, junto con el código y los checkpoints disponibles públicamente.

5. Significado e Impacto

Este trabajo marca un punto de inflexión en la inteligencia artificial aplicada a la fisiología del sueño. Al abordar directamente el problema de la heterogeneidad de los dispositivos y la falta de canales, OSF permite el despliegue de modelos de IA más robustos en entornos clínicos reales (como estudios domiciliarios con sensores limitados). Además, establece las "mejores prácticas" para el pre-entrenamiento de modelos fundamentales en series temporales biomédicas, demostrando que la combinación de datos masivos, arquitecturas grandes y estrategias de invarianza específicas del dominio es la clave para lograr generalización real.

El código y los modelos están disponibles en: https://github.com/yang-ai-lab/OSF-Open-Sleep-FM.