Synthetic Data Generation for Long-Tail Medical Image Classification: A Case Study in Skin Lesions

Riepilogo Tecnico: Generazione di Dati Sintetici per la Classificazione di Immagini Mediche a Coda Lunga

Enunciato del Problema
Le distribuzioni di classe a coda lunga sono pervasive nei dataset medici multiclasse, ponendo sfide significative per i modelli di deep learning. In questi scenari, poche classi dominanti "testa" contengono la maggior parte dei campioni, mentre numerose classi "coda" clinicamente critiche (spesso corrispondenti a malattie rare ad alto rischio) sono gravemente sottorappresentate. I modelli addestrati su tali dati sbilanciati raggiungono tipicamente una precisione complessiva competitiva ma falliscono nel generalizzare alle classi coda. Le soluzioni esistenti, categorizzate in progettazione di funzioni di perdita, miglioramenti dell'architettura del modello e riequilibrio delle classi, offrono solo miglioramenti marginali. Questi metodi spesso dipendono da una meticolosa regolazione degli iperparametri, mostrano una generalizzazione limitata tra modelli diversi e faticano a scalare a causa della natura deterministica dell'augmentation dei dati artigianale. Inoltre, sebbene i modelli di diffusione abbiano dimostrato successo nella generazione di immagini naturali, la loro applicazione alla diagnostica medica specifica di dominio e a coda lunga rimane poco esplorata, in particolare per quanto riguarda la generazione di campioni diversificati ma clinicamente significativi.

Metodologia
Gli autori propongono una pipeline di augmentation di dati sintetici guidata da modelli di diffusione, adattata per la Classificazione di Immagini Mediche a Coda Lunga (LTMIC), dimostrata specificamente sulla classificazione delle lesioni cutanee. L'approccio consta di tre fasi principali:

1. Architettura del Modello di Diffusione per Inpainting:
   Il nucleo del metodo è un innovativo modello di diffusione per inpainting. A differenza dei modelli standard basati su embedding testuale, inadatti a descrizioni mediche altamente specializzate, gli autori addestrano una semplice rete di embedding di classe completamente connessa. Il modello riceve quattro input: l'immagine dermatoscopica originale ($x$), una maschera binaria dell'area di interesse ($b$), lo sfondo dell'immagine dermatoscopica ($d$) e l'etichetta di classe ($y$). L'immagine e lo sfondo vengono codificati in uno spazio latente tramite un Variational Auto Encoder (VAE) pre-addestrato. La maschera binaria viene ridimensionata per corrispondere alle dimensioni latenti. Questi componenti vengono concatenati e immessi in una UNet parametrizzata da $\theta_d$ per stimare il rumore aggiunto durante il processo di diffusione in avanti. Una rete di embedding di classe parametrizzata da $\theta_c$ fornisce la condizionamento dell'etichetta.

2. Fine-Tuning con Selezione dei Campioni:
   Per garantire che il modello di diffusione apprenda la distribuzione di ciascuna classe in modo uniforme, viene adottata una strategia di sovracampionamento durante il fine-tuning. Gli autori sfruttano i risultati di lavori precedenti che suggeriscono che i campioni "approssimativamente nella distribuzione" (AID) (quelli leggermente lontani dal centro della distribuzione) sono i più importanti per l'addestramento di modelli generativi. Di conseguenza, viene utilizzato un metodo di rilevamento delle anomalie mediche auto-supervisionato per calcolare i punteggi di anomalia per ciascun campione. Vengono assegnati pesi in base alla distanza di un campione dalla mediana del punteggio di anomalia della sua classe; i campioni più vicini alla regione AID ricevono pesi più elevati e vengono ripetuti per bilanciare il dataset di fine-tuning tra tutte le classi.

3. Generazione Sintetica e Filtraggio OOD:
   Il modello fine-tuned genera immagini sintetiche utilizzando come input le immagini di sfondo e le etichette di classe del dataset di addestramento originale. A causa della natura stocastica dei modelli di diffusione, il dataset sintetico iniziale ($D_{syn}$) contiene campioni "rumorosi", inclusi immagini Fuori Distribuzione (OOD) che potrebbero non essere clinicamente valide. Per affrontare questo problema, viene applicato un meccanismo di filtraggio OOD medico a $D_{syn}$ per produrre un dataset sintetico pulito ($D_{clean}^{syn}$). Il dataset di addestramento finale è formato combinando il dataset originale a coda lunga ($D_{LT}$) con $D_{clean}^{syn}$. Un iperparametro $\gamma$ controlla la percentuale di campioni puliti mantenuti dal dataset sintetico per ciascuna classe.

Contributi Chiave
Il documento delinea tre contributi principali:

1. Pipeline Innovativa: Una pipeline di generazione di dati sintetici che integra un modello di diffusione per inpainting con filtraggio OOD, progettata specificamente per affrontare il problema LTMIC.
2. Architettura Specializzata: Una nuova architettura di modello di diffusione per inpainting adattata per la generazione di immagini mediche, che utilizza una rete di embedding di classe dedicata e condizionamento dello sfondo invece di embedding testuali generici.
3. Strategia di Campionamento: Una strategia di campionamento basata su OOD per il fine-tuning di modelli di diffusione su dataset medici a coda lunga, che dà priorità ai campioni AID per migliorare la qualità dei dati generati per le classi coda.

Risultati Sperimentali
Il metodo è stato valutato sul dataset di classificazione delle lesioni cutanee ISIC2019, un benchmark contenente 25.331 immagini distribuite su 8 classi con forte sbilanciamento (da 12.875 campioni per Nevus a soli 62 per Lesioni Vascolari). Gli esperimenti hanno utilizzato una validazione incrociata a 5 fold e un semplice classificatore basato su CNN (senza ensemble o dati esterni) per isolare l'impatto del metodo di generazione dei dati.

• Prestazioni Complessive: L'approccio proposto ha raggiunto una Precisione Multiclasse Bilanciata (BMA) di 0,802, superando significativamente i metodi attuali all'avanguardia (SOTA), inclusi ensemble multi-modello (MRE) e altri approcci basati su diffusione (ad es. CBDM, Dreambooth). In particolare, ha superato il metodo MRE di circa il 4% in termini di BMA senza richiedere dati esterni.
• Miglioramento delle Classi Coda: I guadagni più significativi sono stati osservati nelle classi coda. Per la classe con il minor numero di campioni (Dermatofibroma, DF), il metodo ha prodotto un miglioramento della precisione superiore al 28% rispetto alla baseline senza dati sintetici.
• Impatto del Filtraggio OOD: Lo studio ha rilevato che il filtraggio OOD è cruciale per bilanciare le prestazioni. Sebbene i dati sintetici senza filtraggio abbiano migliorato significativamente le classi coda, hanno introdotto rumore che ha leggermente degradato le prestazioni delle classi testa. Il processo di filtraggio ha ripristinato le prestazioni delle classi testa mantenendo i guadagni nelle classi coda.
• Sensibilità agli Iperparametri: Uno studio di ablazione su $\gamma$ (il tasso di ritenzione dei campioni sintetici puliti) ha rivelato che le prestazioni migliorano drasticamente anche con una piccola quantità di dati sintetici, ma si stabilizzano o diminuiscono quando $\gamma$ si avvicina a 1 (nessun filtraggio). Il range ottimale per $\gamma$ è stato identificato tra 0,2 e 0,6.

Significato e Affermazioni
Gli autori affermano che i loro risultati dimostrano l'efficacia dell'augmentation basata su diffusione nel mitigare lo sbilanciamento a coda lunga e nel migliorare la robustezza della classificazione medica. Generando campioni sintetici diversificati, realistici e clinicamente significativi, il metodo affronta la scarsità di dati delle malattie rare senza fare affidamento su ensemble di modelli complessi o funzioni di perdita specializzate sensibili agli iperparametri. Il documento conclude che, sebbene l'approccio attuale trattasse la generazione dei dati e la classificazione come processi indipendenti, il lavoro futuro si concentrerà sull'ottimizzazione congiunta di questi componenti per migliorare ulteriormente le prestazioni LTMIC. Il metodo è presentato come agnostico rispetto al modello, capace di integrazione con qualsiasi architettura di classificazione.