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

Dit paper introduceert DerMAE, een methode die klasse-geconditioneerde diffusiemodellen en MAE-distillatie combineert om een robuust en lichtgewicht Vision Transformer-model te trainen voor de classificatie van huidlaesies, wat de prestaties verbetert ondanks onbalans in de dataset en efficiënte inferentie op mobiele apparaten mogelijk maakt.

De kern

Stel je voor dat je een dokter bent die huidkanker moet diagnosticeren. Het probleem is dat er veel te weinig foto's zijn van de gevaarlijke gevallen (kwaadaardig), terwijl er duizenden foto's zijn van de onschuldige vlekjes (goedaardig).

Als je een computer leert om dit te herkennen met alleen die bestaande foto's, wordt de computer lui. Hij denkt: "Oh, bijna elke vlek is onschuldig, dus ik ga maar zeggen dat alles onschuldig is." Dat is gevaarlijk, want dan mis je de echte kanker.

De auteurs van dit paper (DERMAE) hebben een slimme oplossing bedacht die bestaat uit drie stappen. Laten we het vergelijken met het trainen van een jonge leerling door een meester.

Stap 1: De Kunstenaar die Verzonnen Foto's Maakt (Synthetische Data)

Omdat er te weinig foto's van kanker zijn, hebben ze een digitale kunstenaar ingeschakeld: een Diffusiemodel.

• De Analogie: Stel je voor dat je een schilderij wilt maken van een zeldzame dier, maar je hebt er nog nooit een gezien. In plaats van te wachten tot je er eentje ziet, laat je een kunstenaar (de AI) duizenden nieuwe, realistische schilderijen van dat dier "dromen" en schilderen.
• Wat ze deden: Ze gebruikten een AI om miljoenen nieuwe, verzonnen foto's van huidvlekken te maken. Ze zorgden er specifiek voor dat er evenveel foto's van kanker als van onschuldige vlekken waren. Hierdoor kreeg de computer een eerlijke kans om te leren.

Stap 2: De Meester die Eerst Alles Leert (MAE Pre-training)

Nu hebben ze een enorme hoeveelheid foto's (echt + verzonnen). Ze wilden een heel slimme computer (een groot model genaamd ViT-H) trainen om alles te begrijpen.

• De Analogie: Stel je voor dat je een meesterkunstenaar (de AI) een hele kamer vol met duizenden foto's geeft, maar je bedekt 75% van elke foto met een doekje. De meester moet de ontbrekende stukjes raden en tekenen.
• Het doel: Door te oefenen met het invullen van de ontbrekende stukjes op duizenden foto's, leert de meester niet alleen hoe een vlek eruitziet, maar ook hoe de structuur en patronen van een vlek werken. Hij wordt een expert in het begrijpen van huid, zonder dat hij eerst hoeft te weten welke vlekken kanker zijn. Hij bouwt een sterke basis van kennis.

Stap 3: De Leerling die de Wijsheid Overneemt (Knowledge Distillation)

De meester is nu super slim, maar hij is ook enorm, traag en heeft een enorme computer nodig. Je kunt zo'n zware computer niet in een mobiele telefoon van een huisarts doen.

• De Analogie: Je hebt een slimme, maar zware leraar (de meester) en een jonge, snelle leerling (een klein model voor de telefoon). De leraar kan niet meegaan naar het dorp, maar hij kan wel zijn kennis overdragen.
• De oplossing: Ze laten de jonge leerling kijken naar hoe de meester denkt. De leerling probeert niet alleen het juiste antwoord te geven, maar ook te denken zoals de meester. Zo wordt de kleine, snelle leerling bijna net zo slim als de zware meester, maar past hij wel in je broekzak.

Wat was het resultaat?

Door deze drie stappen te combineren (verzonnen foto's maken + de meester laten oefenen + de leerling slim maken), gebeurde er iets wonderlijks:

1. De computer werd veel beter in het herkennen van kanker, zelfs omdat er weinig echte voorbeelden waren.
2. Het systeem werd zo licht en snel dat het op een mobiele telefoon kan draaien.

Kortom: Ze hebben een slimme manier bedacht om een computer te trainen met "verzonnen" foto's, zodat hij een meester wordt in het herkennen van huidkanker, en ze hebben die wijsheid overgebracht naar een klein programmaatje dat elke arts op zijn telefoon kan gebruiken. Dit helpt om vroege diagnoses te stellen, zelfs in gebieden waar geen specialisten zijn.

Technische samenvatting

Probleemstelling

De classificatie van huidlaesies (huidkanker) kampt met twee fundamentele uitdagingen:

1. Ernstige class-imbalance: Huidkankerdatasets, zoals HAM10000, zijn sterk onbalans. Ongeveer 90% van de samples bestaat uit goedaardige (benigne) laesies, terwijl kwaadaardige (maligne) gevallen sterk ondervertegenwoordigd zijn. Dit leidt tot bias in deep learning-modellen die de meerderheidsklasse prefereren en zeldzame, maar kritieke maligne laesies systematisch verkeerd classificeren.
2. Beperkte data en modelcomplexiteit: Bestaande datasets zijn klein en gevarieerd (door etniciteit, geografie, etc.), wat de prestaties van krachtige modellen zoals Vision Transformers (ViT) beperkt, omdat deze grote hoeveelheden data nodig hebben. Tegelijkertijd vereisen klinische toepassingen (bijv. mobiele apps) lichtgewicht modellen, terwijl de meest accurate ViT-modellen vaak te zwaar zijn voor randapparaten (edge devices).

Methodologie

De auteurs stellen een drie-trapsraamwerk voor (DERMAE) om deze problemen aan te pakken:

1. Synthetische Data Generatie (Conditioned Latent Diffusion)
Om het gebrek aan maligne samples op te lossen, worden synthetische dermatologische afbeeldingen gegenereerd met behulp van een Denoising Diffusion Probabilistic Model (DDPM).

• Architectuur: Een U-Net backbone die een denoising-functie parametriseren.
• Strategie: Er wordt onderscheid gemaakt tussen onvoorwaardelijke generatie (voor de algemene verdeling) en voorwaardelijke generatie (class-conditioned). Bij de voorwaardelijke generatie worden class-embeddings gebruikt tijdens de inferentie om specifiek een gebalanceerd aantal maligne en benigne samples te genereren.
• Training: Het model wordt getraind met een combinatie van Mean Squared Error (MSE) en een perceptuele loss (om visuele trouw te verbeteren en artefacten te verminderen).

2. Zelf-supervisie Pre-training (MAE)
Om de kracht van ViT-modellen te benutten ondanks de beperkte echte data, wordt een Masked Autoencoder (MAE) strategie toegepast.

• Model: Een ViT-Huge (ViT-H/16) model.
• Proces: Het model wordt vooringesteld (pre-trained) op de grote hoeveelheid synthetische data. Tijdens training worden 75% van de image-patches gemaskeerd, en het model moet de ontbrekende delen reconstrueren.
• Doel: Dit dwingt het model om robuuste, semantisch betekenisvolle representaties van laesiestructuren te leren in plaats van oppervlakkige correlaties, wat de generalisatie verbetert.

3. Kennisdistillatie (Knowledge Distillation)
Om het model geschikt te maken voor mobiele apparaten, wordt kennis overgedragen van het zware "leraar"-model naar een lichtgewicht "student"-model.

• Leraar: De geëncodeerde ViT-H (vooringesteld via MAE).
• Student: Een kleinere ViT-Base (ViT-B/16) of EfficientNet-B0.
• Methode: De student wordt getraind om zowel de echte labels (hard targets) als de waarschijnlijkheidsverdeling van de leraar (soft targets) na te bootsen, gebruikmakend van een combinatie van Cross-Entropy en Kullback-Leibler-divergentie.

Belangrijkste Bijdragen

• Integratie van Diffusie en MAE: Een unieke combinatie van class-conditioned diffusiemodellen voor data-augmentatie en MAE voor zelf-supervisie om de prestaties van ViT-modellen op kleine, onbalans datasets te maximaliseren.
• Efficiënte Deployability: Het demonstreren dat kennisdistillatie de superieure representatieve kracht van grote ViT-modellen kan overdragen naar compacte modellen (zoals EfficientNet-B0) zonder significante verlies in nauwkeurigheid, waardoor real-time inferentie op mobiele apparaten mogelijk wordt.
• Verbeterde Kwaliteit: Het gebruik van perceptuele loss en class-conditioning resulteert in synthetische afbeeldingen met hogere visuele kwaliteit en diversiteit, wat essentieel is voor het trainen van betrouwbare diagnostische modellen.

Resultaten

De experimenten zijn uitgevoerd op de HAM10000 dataset (8.012 trainingsamples, 2.003 validatie, 1.512 test).

• Prestatieverbetering: Baseline ViT-modellen presteerden slecht door data-schaarste. Het toepassen van MAE-pretraining op synthetische data leidde tot aanzienlijke verbeteringen.
  • De combinatie van MAE + Synthetische Data bracht de F1-score voor maligne/benigne classificatie van ViT-H van ~0.40 (baseline) naar 0.8715.
• Distillatie Effect: Het distilleren van de ViT-H naar een ViT-B of EfficientNet-B0 behield een groot deel van deze prestaties.
  • De EfficientNet-B0 met distillatie en synthetische data bereikte een F1-score van 0.9010 en een nauwkeurigheid van 0.9051 voor maligne/benigne classificatie. Dit is een aanzienlijke verbetering ten opzichte van de baselines en zelfs beter dan de Derm-t2im benchmark.
• Efficiëntie: Het uiteindelijke studentmodel (EfficientNet-B0) vereist slechts ~0.6 GFLOPs en ~5M parameters, wat het ideaal maakt voor mobiele implementatie, in tegenstelling tot het ViT-H model dat ~55-60 GFLOPs vereist.

Betekenis en Conclusie

Het paper toont aan dat de combinatie van synthetische data-generatie (via diffusie), zelf-supervisie (via MAE) en kennisdistillatie een effectieve oplossing biedt voor de klassieke uitdagingen in medische beeldanalyse: data-imbalance en computational constraints.

De methode stelt artsen in staat om nauwkeurige, mobiele diagnostische tools te gebruiken die goed presteren op zeldzame, kwaadaardige laesies. Dit kan de toegang tot specialistische dermatologische zorg verbeteren, vooral in gebieden waar specialisten schaars zijn, door betrouwbare eerste screenings mogelijk te maken via smartphones.