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

Os autores apresentam o OSF, uma família de modelos de fundação para sono que alcança desempenho state-of-the-art em diversas tarefas de previsão, baseada na curadoria de um vasto corpus de dados, na criação do benchmark SleepBench e na descoberta de que o aprendizado de características invariante a canais e a escalabilidade sistemática são essenciais para a generalização.

O essencial

Imagine que o sono é como uma orquestra gigante tocando dentro do seu corpo. Para entender se a música está boa (você está dormindo bem) ou se há um instrumento desafinado (um problema de saúde), os médicos usam um equipamento chamado Polissonografia (PSG). Esse equipamento conecta vários sensores no seu corpo: na cabeça (cérebro), no peito (respiração), no coração e nas pernas.

O problema é que, na vida real, nem sempre conseguimos usar todos os sensores. Às vezes, o sensor do cérebro cai, às vezes o paciente faz o teste em casa e só tem sensores de respiração. É como tentar ouvir uma sinfonia completa, mas só conseguir ouvir os violinos ou só os tambores.

Até agora, os "cérebros de computador" (Inteligência Artificial) treinados para analisar esse sono funcionavam muito bem quando tinham todos os instrumentos tocando. Mas, se faltasse um sensor, eles ficavam confusos e cometiam muitos erros.

É aqui que entra o OSF (Open Sleep Foundation Model), o novo "maestro" criado pelos pesquisadores da UCLA e Emory University.

Aqui está a história de como eles fizeram isso, explicada de forma simples:

1. O Grande Banco de Dados (SleepBench)

Antes de criar o novo maestro, os pesquisadores precisaram de uma biblioteca de músicas gigantesca para treinar. Eles reuniram 166.500 horas de gravações de sono de 9 fontes diferentes.

• A Analogia: Imagine que eles pegaram milhares de álbuns de música de diferentes gêneros, épocas e gravadoras e misturaram tudo em um único "Spotify do Sono". Isso criou um banco de dados chamado SleepBench, que é aberto para todo mundo usar.

2. O Treinamento: Aprendendo a Ouvir "Sem" Sensores

O grande segredo do OSF não foi apenas ler mais dados, mas mudar como ele aprendeu.

• O Problema Antigo: Os modelos antigos eram como alunos que decoravam a partitura inteira. Se faltasse uma nota (um sensor), eles não sabiam o que fazer.
• A Solução do OSF: Eles treinaram o modelo com uma técnica especial chamada "Mascaramento de Canais".
  • A Analogia: Imagine que você está ensinando alguém a reconhecer uma música de rock. Em vez de deixar a música tocar inteira, você corta o som dos guitarristas em metade das vezes e corta o som da bateria na outra metade. Você força o aluno a aprender a reconhecer a música mesmo sem ouvir todos os instrumentos.
  • Resultado: O OSF aprendeu a entender o sono focando no que é essencial, independentemente de quais sensores estão funcionando. Ele se tornou "invariante aos canais".

3. A Lei da Escala: Quanto Mais, Melhor

Os pesquisadores descobriram que, ao contrário do que se pensava, quanto mais dados e quanto maior o "cérebro" (modelo) do computador, melhor ele ficava.

• A Analogia: É como estudar para uma prova. Se você ler apenas um livro, você passa. Se você ler 100 livros de diferentes autores e tiver um cérebro capaz de guardar tudo isso, você se torna um especialista que consegue responder a qualquer pergunta, mesmo as mais difíceis.
• Eles aumentaram o tamanho do modelo e a quantidade de dados, e o desempenho melhorou consistentemente.

4. O Resultado: O Maestro Perfeito

O OSF foi testado em 9 conjuntos de dados diferentes e em várias tarefas:

• Estágios do Sono: Saber se você está acordado, dormindo leve, dormindo fundo ou sonhando.
• Detecção de Doenças: Identificar apneia (parada da respiração), falta de oxigênio e até prever riscos de doenças cardíacas e diabetes.
• A Grande Vitória: O OSF funcionou tão bem que, mesmo quando faltavam sensores (como se o paciente tivesse feito o teste em casa com equipamento simples), ele ainda acertava mais do que os melhores modelos antigos.

Resumo da Ópera

Os pesquisadores criaram um novo "cérebro" para analisar o sono que é:

1. Robusto: Não se confunde se faltar um sensor.
2. Educativo: Aprendeu com uma quantidade massiva e diversa de dados.
3. Versátil: Serve tanto para saber se você dormiu bem quanto para alertar sobre doenças graves.

Eles liberaram tudo (o código e os dados) para que qualquer cientista ou médico no mundo possa usar essa tecnologia para melhorar a saúde do sono de todos nós. É como se eles tivessem ensinado um computador a ouvir a "música do corpo" perfeitamente, mesmo com ruídos ou instrumentos faltando.

Resumo técnico

Título: OSF: Sobre o Pré-treinamento e Escalonamento de Modelos de Base para Sono (Sleep Foundation Models)

1. Problema e Motivação

O polissonografia (PSG) é o padrão-ouro para avaliação do sono, mas enfrenta desafios significativos na aplicação de modelos de aprendizado de máquina:

• Heterogeneidade: Existem grandes variações nos dispositivos de gravação, protocolos de aquisição e coortes de pacientes.
• Canais Incompletos: Em cenários reais (como estudos domiciliares ou dispositivos vestíveis), canais críticos (ex: EEG, EOG) podem estar ausentes ou sofrer quedas durante a noite.
• Falta de Generalização: Modelos de Fundação (Foundation Models - FMs) existentes para sono, embora promissores, falham em generalizar quando os canais de entrada durante a inferência não correspondem aos usados no pré-treinamento.
• Ausência de Benchmarks Abertos: Não existia um benchmark abrangente e totalmente aberto para avaliar sistematicamente estratégias de pré-treinamento e padrões de escalonamento para modelos de sono.

2. Metodologia

2.1. SleepBench: O Novo Benchmark

Os autores criaram o SleepBench, o maior benchmark aberto e multi-fonte para modelos de sono:

• Dados: 166.500 horas de gravações de PSG de 9 fontes públicas (incluindo SHHS, MESA, CHAT, etc.), abrangendo mais de 21.000 noites de sono.
• Padronização: Utilização de um conjunto compartilhado de 12 canais padronizados (EEG, EOG, Respiração, Cardíaco, Soma) e um pipeline de pré-processamento consistente (normalização z-score, remoção de artefatos, segmentação em épocas de 30 segundos).
• Divisão: Separação rigorosa entre coortes de pré-treinamento (in-domain) e coortes externas para avaliação fora de domínio (OOD).

2.2. Estudo Sistemático de Pré-treinamento

A equipe avaliou quatro famílias de objetivos de aprendizado auto-supervisionado (SSL):

1. Aprendizado Contrastivo: SimCLR.
2. Reconstrução: MAE e VQ-VAE.
3. Modelagem Autoregressiva: Autoregression.
4. Auto-Distilação: DINO.

Eles conduziram experimentos controlados para identificar quais escolhas de design melhoram a transferência e a robustez, focando em três áreas principais:

• Inferência com Canais Faltantes: Simulação de cenários onde canais são mascarados (zero) durante a inferência.
• Aumentação de Dados: Comparação entre time-wise masking (mascaramento temporal) e channel masking (mascaramento de canais).
• Leis de Escalonamento (Scaling Laws): Análise do impacto do aumento do tamanho dos dados, capacidade do modelo e mistura de fontes de dados.

2.3. Desenvolvimento do OSF

Com base nas descobertas, foi introduzido o OSF (Open Sleep Foundation Model). O OSF utiliza a arquitetura DINO como base, mas incorpora um "receituário" (recipe) de pré-treinamento otimizado que inclui:

• Mascaramento de Canais: Força o modelo a aprender representações invariantes aos canais disponíveis.
• Escalonamento: Uso de modelos maiores (ViT-85M) e pré-treinamento em grandes volumes de dados multi-fonte.

3. Contribuições Principais e Descobertas Chave

O artigo estabelece três achados críticos que guiam o desenvolvimento de FMs para sono:

1. Falha na Generalização com Canais Faltantes: Modelos de sono existentes sofrem quedas drásticas de desempenho quando canais críticos (como EEG para estadiamento ou respiração para apneia) estão ausentes na inferência.
2. Aprendizado de Características Invariantes a Canais é Essencial: O pré-treinamento que inclui explicitamente o channel masking (mascarar aleatoriamente 50% dos canais durante o treino) é crucial. Isso ensina o modelo a extrair informações robustas independentemente de quais canais estão presentes, melhorando significativamente a transferência para tarefas downstream.
3. Leis de Escalonamento Consistentes: Ao contrário de modelos anteriores que saturavam com mais dados, o OSF demonstra que aumentar o tamanho da amostra, a capacidade do modelo e a mistura de dados multi-fonte resulta em melhorias consistentes no desempenho downstream.

4. Resultados

O OSF foi avaliado em 9 conjuntos de dados e 8 tarefas downstream (estadiamento do sono, detecção de eventos como apneia/arousal, e predição de doenças).

• Desempenho State-of-the-Art (SOTA): O OSF superou consistentemente todos os métodos comparados (incluindo SleepFM, SimCLR, DINO, MAE) em tarefas de estadiamento do sono e detecção de eventos, tanto em linear probing quanto em fine-tuning completo.
  • Exemplo: No estadiamento do sono (MROS), o OSF alcançou 97.3% de AUC (vs 96.4% do SleepFM) e 90.4% de AUPRC.
• Robustez a Canais Faltantes: O OSF manteve desempenho superior em cenários realistas onde canais de cérebro ou respiração estavam ausentes, demonstrando maior adaptabilidade clínica.
• Eficiência em Few-Shot: O modelo mostrou superioridade significativa em cenários com poucos dados rotulados (1-shot, 5-shot), indicando representações mais ricas e generalizáveis.
• Predição de Doenças: O OSF alcançou os melhores resultados na predição de doenças a nível de paciente (Doença Coronariana, Diabetes, Hipertensão), superando métodos supervisionados e outros modelos de base.
• Escalonamento: Gráficos de escalonamento mostraram que o OSF continua a melhorar com o aumento de dados e tamanho do modelo, enquanto modelos base (como SleepFM original) tendiam a saturar.

5. Significado e Impacto

• Padrão de Referência: O lançamento do SleepBench e do OSF estabelece um novo padrão para pesquisa em modelos de base para sono, permitindo comparações justas e reprodutíveis.
• Viabilidade Clínica: Ao resolver o problema da heterogeneidade de canais e dispositivos, o OSF torna os modelos de IA mais viáveis para uso em ambientes clínicos reais e estudos domiciliares, onde a coleta completa de sinais nem sempre é possível.
• Diretrizes de Design: O trabalho fornece diretrizes práticas para a comunidade de IA em saúde: o uso de channel masking no pré-treinamento e o escalonamento conjunto de dados e modelos são essenciais para construir FMs robustos para sinais fisiológicos.
• Código Aberto: O código e os checkpoints estão disponíveis publicamente, fomentando a colaboração e o avanço rápido na área de saúde digital.

Em resumo, este trabalho não apenas apresenta um modelo superior, mas também desvenda os mecanismos fundamentais (invariância de canais e escalonamento) que permitem que modelos de base para sono sejam generalizáveis e robustos em cenários do mundo real.