DerMAE: Improving skin lesion classification through conditioned latent diffusion and MAE distillation

Il paper DerMAE affronta lo squilibrio delle classi nei dataset di lesioni cutanee combinando la generazione di immagini sintetiche tramite modelli di diffusione condizionati, il preaddestramento auto-supervisionato MAE su grandi modelli ViT e la distillazione della conoscenza per trasferire queste rappresentazioni robuste in modelli leggeri adatti all'inferenza su dispositivi mobili in ambito clinico.

L'essenziale

Immagina di dover insegnare a un medico robot a riconoscere i nei sulla pelle. Il problema è che il "libro di testo" (il dataset di immagini) con cui lo stiamo addestrando è molto sbilanciato: è pieno zeppo di foto di nei benigni (innocui), ma ha pochissime foto di nei maligni (pericolosi).

È come se volessi insegnare a un bambino a riconoscere i leoni mostrandogliene solo uno ogni mille gatti. Il bambino imparerà a pensare che "tutti i felini sono gatti" e non saprà mai riconoscere un leone quando lo vedrà davvero.

Gli autori di questo studio (un gruppo di ricercatori brasiliani) hanno risolto questo problema con una strategia in tre atti, che potremmo chiamare "Il Metodo del Maestro, dello Studente e del Maghetto".

1. Il Maghetto: Creare nuovi amici (Generazione Sintetica)

Prima di tutto, hanno usato un "maghetto" digitale chiamato Diffusione Latente.

• Cos'è? È un'intelligenza artificiale che funziona un po' come un artista che parte da un quadro pieno di "nebbia" e la pulisce poco alla volta fino a rivelare un'immagine.
• Cosa ha fatto? Poiché mancavano foto di nei pericolosi, il maghetto ha "immaginato" e creato migliaia di nuovi nei maligni sintetici. Non sono veri, ma sono così realistici che sembrano veri.
• Il trucco: Hanno insegnato al maghetto a creare esattamente il numero di nei maligni che mancava, così da bilanciare la classe. Ora il "libro di testo" ha un numero uguale di gatti e leoni.

2. Il Maestro: Studiare con un libro gigante (Pre-training MAE)

Ora che abbiamo un libro di testo perfetto (reale + sintetico), serve un insegnante molto intelligente.

• Il Maestro (ViT-Huge): Hanno usato un modello enorme chiamato ViT-Huge. È come un professore universitario che legge tutto il libro di testo (migliaia di immagini) e impara a riconoscere ogni minimo dettaglio, ogni sfumatura di colore e forma.
• Il metodo di studio (MAE): Per imparare davvero bene, il professore ha giocato a un gioco: gli hanno mostrato le immagini coprendo il 75% di esse (come se avesse gli occhi bendati su gran parte del viso). Lui doveva indovinare cosa c'era sotto la benda basandosi solo sui pochi pezzi visibili.
• Il risultato: Questo lo ha costretto a capire la struttura globale del nevo, non solo a memorizzare i pixel. È diventato un esperto assoluto.

3. Lo Studente: Il tirocinante intelligente (Distillazione della Conoscenza)

C'è un problema: il "Professore" (il modello ViT-Huge) è troppo grande e pesante. Non può stare nello smartphone di un medico o in un dispositivo portatile perché consumerebbe troppa batteria e memoria.

• Lo Studente (ViT-Base o EfficientNet): Serve un modello piccolo, leggero e veloce, come un tirocinante che deve lavorare sul campo.
• La Distillazione: Invece di far studiare lo studente dal libro di testo (che è difficile e lento), lo fanno studiare guardando cosa fa il Professore.
  • Immagina che il Professore guardi un nevo e dica: "Questo è maligno, ma guarda come è scuro qui e come è irregolare lì".
  • Lo studente non deve solo imparare la risposta giusta, ma deve imitare il ragionamento del professore.
  • In questo modo, lo studente piccolo eredita la "saggezza" del modello gigante, diventando quasi altrettanto bravo, ma rimanendo leggero come una piuma.

Il Risultato Finale

Grazie a questo sistema:

1. Hanno risolto il problema della scarsità di dati (creando quelli mancanti).
2. Hanno addestrato un modello super intelligente (il Maestro).
3. Hanno trasferito quella intelligenza in un modello piccolo e veloce (lo Studente).

In pratica: Ora possiamo avere un'app sul telefono di un medico di base che, anche in zone remote o con pochi dati a disposizione, riesce a dire con grande precisione se un neo è innocuo o pericoloso, aiutando a salvare vite umane senza bisogno di laboratori costosi o specialisti ovunque.

È come se avessimo creato un "genio" che insegna a un "bambino prodigio" a diventare un medico esperto, usando libri di testo inventati ma perfetti per colmare le lacune della realtà.

Sommario tecnico

1. Il Problema

La classificazione delle lesioni cutanee affronta due sfide principali nei dataset medici attuali:

• Squilibrio delle classi (Class Imbalance): I dataset, come HAM10000, sono fortemente sbilanciati a favore delle lesioni benigne (circa il 90% dei campioni), mentre i casi maligni sono sottorappresentati. Questo porta i modelli di deep learning a sviluppare confini decisionali distorti, favorendo la classe maggioritaria e commettendo errori sistematici nella classificazione delle lesioni maligne.
• Limitazioni dei dati e complessità dei modelli: Le architetture avanzate come i Vision Transformer (ViT) richiedono grandi quantità di dati per apprendere rappresentazioni stabili. Tuttavia, i dataset dermatologici sono piccoli e non standardizzati. Inoltre, i modelli ViT di grandi dimensioni (es. ViT-Huge) sono computazionalmente proibitivi per l'implementazione su dispositivi mobili o in ambienti clinici con risorse limitate, dove sono necessari modelli leggeri.

2. Metodologia

Il paper propone un framework in tre fasi per mitigare lo squilibrio dei dati e abilitare l'inferenza efficiente:

A. Generazione di Dati Sintetici (Latent Diffusion)

Per bilanciare il dataset, gli autori utilizzano un modello di Denoising Diffusion Probabilistic Model (DDPM) condizionato alla classe.

• Strategia: Vengono generate immagini sintetiche dermatologiche per aumentare la rappresentanza delle classi maligne.
• Architettura: Si basa su una U-Net per parametrizzare la funzione di denoising.
• Ottimizzazione: Oltre alla classica perdita MSE (Mean Squared Error), viene introdotta una perdita percettiva per migliorare la fedeltà percettiva e ridurre gli artefatti a bassa frequenza.
• Condizionamento: Durante l'inferenza, vengono introdotti embedding di classe per generare specificamente campioni maligni o benigni, garantendo un bilanciamento esplicito.

B. Pre-addestramento Self-Supervised (MAE)

Per sfruttare i dati sintetici generati e apprendere rappresentazioni robuste, viene utilizzato un Masked Autoencoder (MAE).

• Modello Maestro: Viene addestrato un ViT-Huge (ViT-H/16) su un vasto insieme di dati sintetici (e reali).
• Meccanismo: Durante il training, il 75% dei patch dell'immagine viene mascherato. Il modello impara a ricostruire i patch mancanti, costringendolo a catturare strutture globali delle lesioni e rappresentazioni latenti semanticamente significative, piuttosto che semplici correlazioni a basso livello.
• Obiettivo: Questo passaggio permette al modello di apprendere caratteristiche generalizzabili nonostante la scarsità di dati reali maligni.

C. Distillazione della Conoscenza (Knowledge Distillation)

Per rendere il sistema deployabile su dispositivi mobili, le conoscenze apprese dal grande modello ViT-H (Teacher) vengono trasferite a un modello più piccolo (Student).

• Architettura Student: Viene utilizzato un ViT-Base (ViT-B/16) o EfficientNet-B0.
• Strategia: Si utilizza una distillazione a "soft-targets". Lo studente viene addestrato non solo sulle etichette reali (ground-truth), ma anche sulla distribuzione di probabilità del modello Teacher (ViT-H pre-addestrato con MAE).
• Funzione di Loss: La perdita totale combina la Cross-Entropy con le etichette vere e la Divergenza di Kullback-Leibler (KL) tra le distribuzioni dello studente e del Teacher.

3. Risultati Chiave

Gli esperimenti sono stati condotti sul dataset HAM10000, confrontando diverse configurazioni (Baseline, MAE + Sintetico, Distillazione + Augmentation).

• Prestazioni Superiori: L'approccio combinato (Pre-addestramento MAE su dati sintetici + Distillazione) ha ottenuto le migliori prestazioni.
  • Il modello MAE + ViT-H + Synth ha raggiunto un'accuratezza del 89.15% e un F1-score del 89.11% nella classificazione categorica, superando significativamente i baseline ViT e EfficientNet addestrati solo su dati reali.
  • Nella classificazione binaria (Maligno/Benigno), il modello MAE ViT-H distill to EfficientNet-B0 + synth ha ottenuto un'accuratezza del 90.51% e un F1-score del 90.10%.
• Beneficio per le Classi Minoritarie: Le metriche basate sull'F1-score sono migliorate drasticamente, indicando che l'inclusione di dati sintetici ha permesso al modello di apprendere meglio le caratteristiche delle lesioni maligne, riducendo i falsi negativi.
• Efficienza Computazionale:
  • Il modello Teacher (ViT-H) richiede circa 2.5GB di memoria e 55-60 GFLOPs, rendendolo inadatto ai dispositivi edge.
  • Il modello Student distillato (EfficientNet-B0) richiede solo 0.6 GFLOPs e circa 5 milioni di parametri, permettendo un'inferenza efficiente su dispositivi mobili senza sacrificare significativamente l'accuratezza.

4. Contributi Principali

1. Framework Ibrido: Integrazione innovativa di generazione sintetica condizionata (Diffusion), pre-addestramento self-supervised (MAE) e distillazione della conoscenza per affrontare simultaneamente lo squilibrio dei dati e i vincoli computazionali.
2. Validazione dell'Uso di Dati Sintetici: Dimostrazione che il pre-addestramento su dati sintetici generati da modelli di diffusione può migliorare le rappresentazioni di modelli ViT su dataset dermatologici reali e sbilanciati.
3. Deployabilità Clinica: Sviluppo di un modello leggero (studente) capace di mantenere alte prestazioni diagnostiche, rendendo possibile l'uso di sistemi di supporto decisionale basati su AI direttamente su dispositivi mobili in contesti clinici reali (point-of-care).

5. Significato e Impatto

Questo lavoro è significativo perché offre una soluzione pratica a due problemi critici nella dermatologia computazionale: la scarsità di dati maligni e la necessità di modelli leggeri.

• Clinico: Migliora l'accuratezza diagnostica per le lesioni maligne, riducendo il rischio di diagnosi errate dovute allo squilibrio dei dati.
• Tecnologico: Dimostra che è possibile combinare modelli generativi avanzati (Diffusion) e tecniche di compressione (Distillazione) per creare sistemi medici AI robusti, scalabili e pronti per l'uso su dispositivi con risorse limitate, facilitando l'accesso a diagnosi specializzate in aree remote o con carenza di dermatologi.