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

Il paper introduce OSF, una famiglia di modelli fondazione per il sonno ad alte prestazioni ottenuta attraverso la creazione del benchmark SleepBench e l'ottimizzazione delle strategie di pre-addestramento e scalatura su un vasto corpus di dati eterogenei.

L'essenziale

Immagina il sonno come un orchestra sinfonica. Ogni strumento (i canali del poligrafo: cervello, respiro, cuore, muscoli) suona una nota diversa. Per capire davvero come sta suonando l'orchestra (la salute del paziente), dovresti ascoltare tutti gli strumenti insieme.

Il problema? Nella vita reale, spesso alcuni strumenti si rompono o mancano. A volte non abbiamo il microfono del cervello (EEG), a volte manca quello del respiro, e a volte i sensori si staccano durante la notte. I vecchi "esperti" di intelligenza artificiale (i modelli precedenti) erano come un direttore d'orchestra che, se mancava anche solo un violino, smetteva di capire la musica o suonava male.

Questo paper presenta OSF, un nuovo "Direttore d'Orchestra" super-intelligente che ha imparato a suonare anche quando mancano metà degli strumenti.

Ecco come ci sono riusciti, passo dopo passo:

1. La Grande Biblioteca del Sonno (SleepBench)

Prima di insegnare a un bambino a leggere, devi dargli molti libri. I ricercatori hanno raccolto una quantità enorme di dati: 166.500 ore di registrazioni del sonno da 9 fonti diverse.

• L'analogia: Immagina di avere una biblioteca con milioni di diari di sonno di persone diverse, di tutte le età e con diversi tipi di problemi. Hanno creato un "campo di addestramento" chiamato SleepBench dove l'IA può studiare senza essere influenzata da un solo tipo di paziente.

2. Tre Scoperte Fondamentali (Le Regole d'Oro)

Analizzando come i vecchi modelli fallivano, ne hanno scoperti tre di errori e come correggerli:

• Errore 1: La fragilità. Se togli un canale (es. il respiro), i vecchi modelli andavano in tilt.
  • La soluzione: L'IA deve imparare che il "significato" del sonno non dipende da quale strumento suona, ma da come suonano tutti insieme. Deve essere indipendente dagli strumenti.
• Errore 2: L'allenamento sbagliato. I vecchi modelli venivano allenati solo con dati perfetti.
  • La soluzione: Hanno iniziato ad allenare l'IA togliendo a caso i canali durante l'addestramento (come un musicista che si allena suonando con gli occhi bendati o senza uno strumento). Così, quando arriva il momento della verità e manca un dato, l'IA non va in panico.
• Errore 3: Non basta poco. I modelli precedenti si fermavano presto, anche se davamo loro più dati.
  • La soluzione: Hanno scoperto che più dati, modelli più grandi e più varietà di fonti (mescolare dati di ospedali diversi) rendono l'IA sempre più brava. È come dire: "Più libri leggi, più diventi saggio".

3. OSF: Il Nuovo Campione

Grazie a queste regole, hanno creato OSF (Open Sleep Foundation Model).

• Cosa fa? È un modello che può prevedere se una persona ha l'apnea, il diabete o l'ipertensione analizzando il sonno.
• Perché è speciale?
  • È robusto: Se manca il segnale del cervello, usa il respiro e il cuore per capire cosa sta succedendo.
  • È efficiente: Impara molto velocemente, anche con pochi esempi (come un bambino che impara a leggere dopo aver visto poche parole).
  • È preciso: Ha battuto tutti i record precedenti su 9 dataset diversi, sia per classificare le fasi del sonno (svegli, REM, sonno profondo) che per diagnosticare malattie.

4. Il Segreto Finale: Come leggere la partitura

Hanno scoperto anche un trucco per leggere la "partitura" completa della notte. Invece di fare una media semplice di tutto il sonno (come mescolare tutte le note in un'unica zuppa), l'IA di OSF sa selezionare i momenti chiave.

• L'analogia: Se vuoi sapere se una persona ha l'apnea, non devi ascoltare tutto il sonno, ma devi concentrarti sui momenti in cui il respiro si ferma. OSF sa esattamente quali "momenti" della notte sono importanti per la malattia e quali no, aggregando le informazioni in modo intelligente.

In Sintesi

Questo paper ci dice che per creare un'intelligenza artificiale che capisce il sonno (e le malattie correlate), non basta avere tanti dati. Bisogna:

1. Allenarla con dati vari e misti.
2. Insegnarle a funzionare anche quando mancano pezzi dell'informazione.
3. Farla diventare grande e potente (scalare).

Il risultato è un modello che può essere usato anche con dispositivi domestici semplici (che hanno meno sensori) per diagnosticare problemi seri, rendendo la medicina del sonno più accessibile a tutti.

Sommario tecnico

1. Il Problema

La polisonnografia (PSG) è lo standard aureo per la valutazione del sonno, ma l'analisi automatizzata dei dati di sonno affronta sfide significative:

• Eterogeneità dei dati: I segnali variano notevolmente tra diversi dispositivi, coorti di pazienti e protocolli di registrazione.
• Canali mancanti: In scenari reali (es. studi domiciliari o dispositivi indossabili), i canali critici (come EEG o EOG) possono essere assenti o i sensori possono disallinearsi durante la notte. I modelli esistenti spesso falliscono quando si trovano a dover inferire su dati con canali mancanti.
• Mancanza di comprensione dello scaling: Non è chiaro quali scelte di pre-addestramento (pre-training) e strategie di scalatura (aumento di dati e capacità del modello) portino effettivamente a modelli fondazione (Foundation Models - FM) più generalizzabili per la fisiologia del sonno.
• Benchmark frammentati: Mancava un benchmark aperto, completo e su larga scala per valutare sistematicamente diversi obiettivi di apprendimento auto-supervisionato (SSL) sul sonno.

2. Metodologia e Infrastruttura

SleepBench: Il Benchmark

Gli autori hanno creato SleepBench, il più grande benchmark open-source per il sonno, aggregando 166.500 ore di registrazioni PSG da 9 dataset pubblici (inclusi SHHS, MESA, CHAT, ecc.), coprendo oltre 21.000 notti di studio.

• Standardizzazione: I dati sono stati processati con una pipeline standardizzata, normalizzati per notte (z-score) e segmentati in epoche non sovrapposte di 30 secondi.
• Canali: Viene utilizzato un set standardizzato di 12 canali raggruppati per modalità fisiologica: Cerebro (EEG/EOG), Respirazione, Cardiaco (ECG) e Somatico (EMG).
• Split: I dati sono divisi in una coorte interna per il pre-addestramento e una coorte esterna (Out-of-Domain, OOD) per la valutazione rigorosa.

Strategie di Pre-addestramento

Lo studio ha valutato sistematicamente quattro famiglie di obiettivi SSL:

1. Apprendimento Contrastivo: SimCLR.
2. Metodi basati su Ricostruzione: MAE e VQ-VAE.
3. Modellazione Autoregressiva: Autoregression.
4. Auto-distillazione: DINO.

Scoperte Chiave e Progettazione di OSF

Attraverso esperimenti controllati, gli autori hanno identificato tre principi fondamentali che hanno guidato la creazione di OSF (Open Sleep Foundation Model):

1. Fallimento con canali mancanti: I modelli FM esistenti non generalizzano bene quando i canali di input sono parziali durante l'inferenza.
2. Apprendimento di caratteristiche invariate ai canali: È essenziale addestrare il modello a imparare rappresentazioni che siano invarianti rispetto al sottoinsieme di canali disponibili. Questo è stato ottenuto introducendo il Channel Masking (mascheramento casuale del 50% dei canali) durante l'aumento dei dati (data augmentation), combinato con il mascheramento temporale.
3. Legge di Scalatura (Scaling Laws): Aumentare la dimensione del campione di addestramento, la capacità del modello e la miscela di dati multi-sorgente migliora costantemente le prestazioni, a patto di utilizzare un'architettura SSL adatta (come quella con mascheramento dei canali).

OSF è la famiglia di modelli risultante, basata su DINO, che incorpora queste ricette di pre-addestramento e scalatura.

3. Risultati Principali

• Prestazioni State-of-the-Art (SOTA): OSF ha ottenuto le migliori prestazioni su 9 dataset diversi per compiti di livello epoca (staging del sonno, rilevamento di eventi come apnee e arousal) e compiti a livello paziente (predizione di malattie come coronaropatia, diabete e ipertensione).
  • Esempio: Nel task di staging del sonno (linear probing), OSF ha raggiunto un AUC di 97.3 su MROS, superando SleepFM (96.4) e altri metodi SSL.
• Robustezza ai Canali Mancanti: OSF dimostra una resilienza superiore quando canali critici (es. EEG o respiratori) sono assenti durante l'inferenza, mantenendo prestazioni elevate dove i modelli precedenti crollano.
• Efficienza nel Few-Shot: OSF eccelle negli scenari con pochi dati etichettati (1-shot, 5-shot), superando significativamente i metodi concorrenti, grazie alla qualità delle rappresentazioni apprese.
• Scalatura: È stato dimostrato che le prestazioni migliorano linearmente all'aumentare della dimensione del modello (da 1M a 85M parametri) e della quantità di dati di pre-addestramento, specialmente quando si utilizza la miscela multi-sorgente.
• Aggregazione Gerarchica: Per la predizione a livello paziente, l'uso di una selezione "Top-k" delle epoche più rilevanti per la malattia ha funzionato meglio della semplice media pooling, suggerendo che i segnali di malattia sono concentrati in specifici periodi del sonno.

4. Contributi Chiave

1. SleepBench: La creazione di un corpus di dati massivo, diversificato e completamente open-source per la ricerca sui modelli fondazione del sonno.
2. Analisi Sistematica: La prima valutazione controllata che identifica il mascheramento dei canali (channel masking) come strategia cruciale per la robustezza e la generalizzazione nei modelli di sonno.
3. OSF (Open Sleep Foundation Model): Un nuovo modello fondazione che stabilisce un nuovo standard di riferimento (SOTA) su una vasta gamma di task clinici e fisiologici.
4. Linee Guida di Scalatura: La dimostrazione che le leggi di scalatura osservate in altri domini (come il linguaggio o la visione) si applicano anche al sonno, ma richiedono un'architettura SSL specifica (invarianza ai canali) per essere efficaci.

5. Significato e Impatto

Questo lavoro rappresenta un passo fondamentale verso l'uso clinico affidabile dell'IA nel sonno. Dimostrando che i modelli possono essere resi robusti alla mancanza di canali (un problema comune nei dispositivi indossabili e domiciliari) e che la scalatura dei dati porta a miglioramenti reali, OSF abilita:

• Diagnosi più accessibili: Possibilità di utilizzare dispositivi con meno sensori senza sacrificare l'accuratezza diagnostica.
• Generalizzazione: Modelli che funzionano bene su popolazioni diverse e con dispositivi diversi, riducendo il bias legato al dataset di addestramento.
• Riproducibilità: La disponibilità di codice, checkpoint e del benchmark SleepBench favorisce lo sviluppo futuro di modelli fondazione per la fisiologia umana.

In sintesi, OSF non è solo un modello migliore, ma fornisce una "ricetta" chiara su come costruire modelli fondazione per il sonno che siano robusti, scalabili e pronti per il mondo reale.