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

Technische Samenvatting: Synthetische Data-Generatie voor Langstaart Classificatie van Medische Afbeeldingen

Probleemstelling
Langstaart-verdelingen van klassen komen veel voor in multi-class medische datasets en vormen aanzienlijke uitdagingen voor deep learning-modellen. In deze scenario's bevatten enkele dominante "head"-klassen het merendeel van de samples, terwijl talrijke klinisch kritische "tail"-klassen (vaak corresponderend met zeldzame, hoog-risico ziekten) ernstig ondervertegenwoordigd zijn. Modellen die zijn getraind op dergelijk scheefgetrokken data bereiken doorgaans een concurrerende algehele nauwkeurigheid, maar falen in het generaliseren naar tail-klassen. Bestaande oplossingen – gecategoriseerd in ontwerp van verliesfuncties, verbeteringen van modelarchitecturen en herbalancering van klassen – bieden slechts marginale verbeteringen. Deze methoden vertrouwen vaak op zorgvuldige afstelling van hyperparameters, vertonen beperkte generalisatie over verschillende modellen en hebben moeite om te schalen vanwege de deterministische aard van handgemaakte data-augmentatie. Bovendien, hoewel diffusiemodellen succes hebben bewezen in het genereren van natuurlijke afbeeldingen, blijft hun toepassing op domeinspecifieke, langstaart medische beeldvorming onderbelicht, met name wat betreft het genereren van diverse en toch klinisch zinvolle samples.

Methodologie
De auteurs stellen een door diffusiemodellen aangedreven pipeline voor synthetische data-augmentatie voor die is toegespitst op Langstaart Classificatie van Medische Afbeeldingen (LTMIC), specifiek gedemonstreerd op classificatie van huidlaesies. De aanpak bestaat uit drie hoofdfasen:

1. Architectuur van het Inpainting Diffusiemodel:
   De kern van de methode is een nieuw inpainting diffusiemodel. In tegenstelling tot standaard tekst-embedding-modellen, die ongeschikt zijn voor highly gespecialiseerde medische beschrijvingen, trainen de auteurs een eenvoudig volledig verbonden class-embeddingnetwerk. Het model neemt vier invoeren: de originele dermoscopie-afbeelding ($x$), een binaire maskering van het gebied van belang ($b$), de achtergrond van de dermoscopie-afbeelding ($d$), en het class-label ($y$). De afbeelding en achtergrond worden via een vooraf getrainde Variational Auto Encoder (VAE) gecodeerd naar de latente ruimte. Het binaire masker wordt herschaald om overeen te komen met de latente dimensies. Deze componenten worden geconcateneerd en ingevoerd in een UNet geparametriseerd door $\theta_d$ om de ruis te schatten die tijdens het voorwaartse diffusieproces is toegevoegd. Een class-embeddingnetwerk geparametriseerd door $\theta_c$ zorgt voor de label-conditioning.

2. Finetuning met Selectie van Samples:
   Om ervoor te zorgen dat het diffusiemodel de verdeling van elke klas even goed leert, wordt tijdens de finetuning een oversampling-strategie toegepast. De auteurs maken gebruik van bevindingen uit eerder werk die suggereren dat "ongeveer-in-de-verdeling" (AID) samples (die iets verwijderd zijn van het centrum van de verdeling) het belangrijkst zijn voor het trainen van generatieve modellen. Derhalve wordt een zelftoezichtende methode voor medische anomaliedetectie gebruikt om anomalie-scores te berekenen voor elke sample. Gewichten worden toegewezen op basis van de afstand van een sample tot de mediaan van de anomalie-score van zijn klas; samples die dichter bij het AID-gebied liggen, krijgen hogere gewichten en worden herhaald om de finetuning-dataset over alle klassen in evenwicht te brengen.

3. Synthetische Generatie en OOD-filtering:
   Het gefinetunte model genereert synthetische afbeeldingen door de achtergrondafbeeldingen en class-labels van de originele trainingsset als invoer te nemen. Vanwege het stochastische karakter van diffusiemodellen bevat de initiële synthetische dataset ($D_{syn}$) "ruisachtige" samples, waaronder Out-of-Distribution (OOD) afbeeldingen die mogelijk niet klinisch valide zijn. Om dit aan te pakken, wordt een medisch OOD-filtermechanisme toegepast op $D_{syn}$ om een schone synthetische dataset ($D_{clean}^{syn}$) te produceren. De definitieve trainingsdataset wordt gevormd door de originele langstaart-dataset ($D_{LT}$) te combineren met $D_{clean}^{syn}$. Een hyperparameter $\gamma$ regelt het percentage schone samples dat uit de synthetische set per klas wordt behouden.

Belangrijkste Bijdragen
Het artikel schetst drie hoofdbijdragen:

1. Nieuwe Pipeline: Een pipeline voor synthetische data-generatie die een inpainting diffusiemodel integreert met OOD-filtering, specifiek ontworpen om het LTMIC-probleem aan te pakken.
2. Gespecialiseerde Architectuur: Een nieuwe architectuur voor een inpainting diffusiemodel die is toegespitst op generatie van medische afbeeldingen, waarbij gebruik wordt gemaakt van een dedicated class-embeddingnetwerk en achtergrond-conditioning in plaats van generieke tekst-embeddings.
3. Sampling-strategie: Een nieuwe op OOD gebaseerde sampling-strategie voor het finetunen van diffusiemodellen op medische langstaart-datasets, die prioriteit geeft aan AID-samples om de kwaliteit van gegenereerde data voor tail-klassen te verbeteren.

Experimentele Resultaten
De methode werd geëvalueerd op de ISIC2019 dataset voor classificatie van huidlaesies, een benchmark met 25.331 afbeeldingen over 8 klassen met ernstige onbalans (variërend van 12.875 samples voor Nevus tot slechts 62 voor Vascular Lesion). De experimenten maakten gebruik van 5-voudige cross-validatie en een eenvoudige CNN-gebaseerde classifier (zonder ensembling of externe data) om de impact van de methode voor data-generatie te isoleren.

• Algemene Prestaties: De voorgestelde aanpak behaalde een Balanced Multiclass Accuracy (BMA) van 0,802, wat significant beter presteerde dan huidige State-of-the-Art (SOTA) methoden, waaronder multi-model ensembles (MRE) en andere op diffusie gebaseerde benaderingen (bijv. CBDM, Dreambooth). Opmerkelijk is dat het de MRE-methode met ongeveer 4% in BMA overtrof zonder externe data te vereisen.
• Verbetering van Tail-klassen: De grootste winst werd waargenomen in tail-klassen. Voor de klas met het minste aantal samples (Dermatofibroma, DF) leverde de methode een verbetering van de nauwkeurigheid van meer dan 28% op ten opzichte van de baseline zonder synthetische data.
• Impact van OOD-filtering: De studie vond dat OOD-filtering cruciaal is voor het in evenwicht brengen van de prestaties. Hoewel synthetische data zonder filtering de tail-klassen significant verbeterde, introduceerde het ruis die de prestaties van head-klassen licht verslechterde. Het filterproces herstelde de prestaties van head-klassen terwijl de winst in tail-klassen behouden bleef.
• Gevoeligheid voor Hyperparameters: Een ablatiestudie naar $\gamma$ (het retentiepercentage van schone synthetische samples) onthulde dat de prestaties sterk verbeteren met zelfs een klein aantal synthetische data, maar verzadigen of afnemen naarmate $\gamma$ naar 1 nadert (geen filtering). Het optimale bereik voor $\gamma$ werd geïdentificeerd tussen 0,2 en 0,6.

Betekenis en Claims
De auteurs claimen dat hun resultaten de effectiviteit van op diffusie gebaseerde augmentatie aantonen in het mitigeren van langstaart-onbalans en het versterken van de robuustheid van medische classificatie. Door diverse, realistische en klinisch zinvolle synthetische samples te genereren, pakt de methode de data-schaarste van zeldzame ziekten aan zonder te vertrouwen op complexe model-ensembles of gespecialiseerde verliesfuncties die gevoelig zijn voor hyperparameters. Het artikel concludeert dat, hoewel de huidige aanpak data-generatie en classificatie behandelt als onafhankelijke processen, toekomstig werk zich zal richten op de gezamenlijke optimalisatie van deze componenten om de LTMIC-prestaties verder te verbeteren. De methode wordt gepresenteerd als model-agnostisch, capable van integratie met elke classificatie-architectuur.