Effect of Patch Size on Fine-Tuning Vision Transformers in Two-Dimensional and Three-Dimensional Medical Image Classification

Deze studie toont aan dat het gebruik van kleinere patch-groottes (1, 2 en 4) bij het fine-tunen van Vision Transformers voor 2D- en 3D-medische beeldclassificatie leidt tot aanzienlijk betere prestaties vergeleken met grotere patches, hoewel dit ten koste gaat van de rekenkracht.

De kern

Stel je voor dat je een gigantische, complexe medische scan (zoals een MRI of een röntgenfoto) moet analyseren om te zien of er iets mis is. Een kunstmatige intelligentie (AI) moet dit ook doen. In het verleden gebruikten AI-modellen vaak een "zoom-in" aanpak: ze keken naar kleine stukjes van de foto tegelijk, alsof ze door een vergrootglas keken.

Maar nu is er een nieuwere, krachtigere AI-technologie genaamd Vision Transformer (ViT). Deze werkt niet met een vergrootglas, maar met een mosaïek.

Het Mosaïek-probleem: Hoe groot moeten de tegels zijn?

Om een ViT een foto te laten "lezen", moet de AI de foto eerst in stukjes hakken, net als een mozaïek van tegels. Elke tegel wordt een stukje informatie (een "token") voor de computer.

De vraag die deze onderzoekers zich stelden, was simpel maar cruciaal: Hoe groot moeten die tegels zijn?

• Grote tegels (bijv. 28x28): Je hebt maar een paar tegels nodig om de hele foto te dekken. Het is alsof je een landschap beschrijft met slechts drie grote blokken: "links is groen, rechts is blauw, midden is geel." Je ziet de grote lijnen, maar je mist de details. Een boom wordt dan misschien verward met een struik.
• Kleine tegels (bijv. 1x1): Je hebt duizenden kleine tegels nodig. Nu zie je elk blaadje, elke tak en elke steen. Je hebt een heel gedetailleerd beeld.

Wat hebben de onderzoekers ontdekt?

De onderzoekers hebben dit getest op 12 verschillende medische datasets (zowel platte 2D-foto's als 3D-volumetrische scans, zoals CT-scan blokken). Ze hebben de AI getraind met verschillende "tegelgroottes" en gekeken hoe goed de diagnose was.

Hier zijn de belangrijkste bevindingen, vertaald naar alledaags taal:

1. Kleiner is beter (voor de diagnose)
Het verrassende nieuws is: kleinere tegels werken veel beter.
Of het nu gaat om een 2D-foto of een 3D-scan, de AI maakte veel minder fouten als ze de foto in heel kleine stukjes hakte.

• De analogie: Stel je voor dat je een schilderij moet kopiëren. Als je het doet met grote kwasten (grote tegels), krijg je een vaag beeld. Gebruik je een heel fijn penseel (kleine tegels), dan krijg je een perfect kopie. In de medische wereld betekent dit: de AI ziet kleine afwijkingen (zoals een klein knobbeltje in een long) die bij grote tegels gewoon "weggemorst" werden.

2. Het prijskaartje: Rekenkracht
Maar er is een addertje onder het gras.

• De analogie: Als je een mozaïek maakt met kleine tegels, moet je veel meer tegels plakken. Dat kost veel meer tijd en energie.
• In de computerwereld betekent dit: hoe kleiner de tegels, hoe meer rekenkracht er nodig is. Als je de tegelgrootte halveert, moet de computer soms 16 keer (bij 2D) of zelfs 64 keer (bij 3D) meer rekenwerk verrichten.
• De onderzoekers vonden dat ze voor de beste resultaten (kleinste tegels) wel meer rekenkracht nodig hadden, maar dat het nog steeds mogelijk was op één enkele moderne videokaart (een GPU), wat een groot voordeel is voor ziekenhuizen die niet over supercomputers beschikken.

3. De "Super-Team" strategie
De onderzoekers probeerden ook iets slims: ze lieten drie AI's werken (één met kleine tegels, één met middelgrote, één met heel kleine) en lieten ze hun antwoorden samenvoegen.

• De analogie: Het is alsof je drie specialisten vraagt om naar een patiënt te kijken. De een kijkt naar het grote plaatje, de ander naar de details. Als ze het samen eens zijn, is de diagnose vaak nog nauwkeuriger. Dit "teamwerk" gaf in veel gevallen de allerbeste resultaten.

Waarom is dit belangrijk?

Vroeger dachten veel mensen dat de standaardgrootte van die tegels (vaak 14x14 of 16x16) wel goed genoeg was. Deze studie zegt: Nee, dat is niet genoeg voor medische scans.

• Voor 2D-foto's (zoals huidfoto's of röntgenfoto's) verbeterde de nauwkeurigheid met wel 12% door kleinere tegels te gebruiken.
• Voor 3D-scans (zoals MRI's van het hele lichaam) was de winst nog groter: tot 23%.

Conclusie

De boodschap is simpel: Als je een AI wilt gebruiken om medische scans te analyseren, moet je de "zoom" zo ver mogelijk inzoomen (kleine tegels gebruiken). Je ziet dan de fijne details die levens kunnen redden. Ja, het kost de computer iets meer werk, maar de verbetering in diagnose is het meer dan waard.

Het is alsof je van een wazige foto overstapt naar een 8K-beeld: je ziet plotseling dingen die je eerder helemaal niet kon zien. En dankzij deze studie weten we nu precies hoe we die "lens" moeten instellen voor de beste medische AI.

Technische samenvatting

Probleemstelling

Vision Transformers (ViTs) en hun varianten zijn de state-of-the-art geworden in veel computer vision-taken en worden veel gebruikt als ruggegraat in grote visuele en visueel-taalmodellen. Hoewel er veel onderzoek is gedaan naar architecturale verbeteringen, blijft de impact van patch size (de grootte van de beeldfragmenten) onderbelicht. Dit is een cruciale ontwerpkeuze, omdat de patch-grootte direct bepaalt hoeveel tokens (fragmenten) er worden gegenereerd en dus hoeveel fijnmazige ruimtelijke informatie het model kan vastleggen.

Bestaande literatuur heeft dit onderwerp vooral onderzocht op natuurlijke afbeeldingen (2D), vaak met modellen die van scratch worden getraind. Er is echter een duidelijke kennislacune wat betreft:

1. De systematische impact van patch-grootte op medische beeldvorming (zowel 2D als 3D).
2. De prestaties bij fine-tuning (de standaardbenadering in de medische domeinen vanwege beperkte datasets) in plaats van training van scratch.
3. De haalbaarheid van dergelijke analyses met beperkte rekenkracht (één GPU), aangezien kleinere patches de computationele kosten exponentieel verhogen.

Methodologie

De auteurs hebben een uitgebreide evaluatie uitgevoerd om te bepalen hoe variaties in patch-grootte de classificatieprestaties van ViTs beïnvloeden in medische contexten.

• Datasets: Er zijn 12 datasets gebruikt uit de MedMNIST V2 collectie:
  • 7 2D-datasets: Inclusief modaliteiten zoals ultrasone geluiden, dermatoscopie, OCT, X-ray en microscopie.
  • 5 3D-datasets: Inclusief CT-scans en MRI (bijv. nier, fracturen, knobbels).
  • Om het experiment op één GPU uitvoerbaar te houden, zijn alle beelden herschaald naar een kleine resolutie: 28x28 voor 2D en 28x28x28 voor 3D.
• Model: Er is gebruik gemaakt van een ViT-Small model (22 miljoen parameters), vooraf getraind op ImageNet.
  • Voor 2D-data werd de classificatie-head vervangen.
  • Voor 3D-data werden de 2D-vooraf getrainde gewichten uitgebreid naar 3D-kernen via "weight inflation" en werden positiële embeddings geïnterpoleerd.
• Experimenteel Ontwerp:
  • Er zijn patch-groottes getest van 1, 2, 4, 7, 14 en 28.
  • De modellen werden fine-tuned met de AdamW optimizer gedurende 80 epochs.
  • Er werd een ensemble-strategie toegepast waarbij de voorspellingen van modellen met patch-groottes 1, 2 en 4 werden gefuseerd (gemiddeld) om te zien of dit de prestaties verder verbeterde.
• Computationele Analyse: De auteurs analyseren de wiskundige relatie tussen patch-grootte en rekentijd. Omdat de complexiteit van self-attention kwadratisch groeit met het aantal tokens ($T^2$), leidt het halveren van de patch-grootte tot een $N^4$-toename in kosten voor 2D en een $N^6$-toename voor 3D.

Belangrijkste Bijdragen

1. Eerste systematische evaluatie: Dit is de eerste studie die de effecten van progressief kleinere patch-groottes grondig evalueert voor zowel 2D als 3D medische beeldclassificatie met ViTs.
2. Fine-tuning focus: In tegenstelling tot eerdere studies die vaak van scratch trainen, focust deze studie op fine-tuning, wat praktischer is voor medische toepassingen met beperkte data.
3. Haalbaarheid op één GPU: De studie demonstreert dat gedetailleerde tokenisatie-analyses mogelijk zijn op een enkele, bescheiden GPU (NVIDIA RTX 4090) door gebruik te maken van geschikte dataset-resoluties, waardoor de drempel voor onderzoekers wordt verlaagd.
4. Open Source: De volledige implementatie is publiek beschikbaar gesteld op GitHub om transparantie en reproduceerbaarheid te waarborgen.

Resultaten

De resultaten tonen een duidelijke trend: kleinere patch-groottes leiden tot betere classificatieprestaties, maar ten koste van hogere rekenkosten.

• 2D Datasets:
  • Kleinere patches (1, 2 en 4) presteerden consistent beter dan de standaard grote patches (zoals 14 of 28).
  • Patch-grootte 2 leverde over het algemeen de beste individuele resultaten op.
  • De ensemble-methode (fuseren van patches 1, 2 en 4) boekte de beste resultaten in bijna alle gevallen.
  • Verbetering: Er werd een toename in Balanced Accuracy gezien tot 12,78% (bijv. patch 2 vs. 28 op de OrganMNIST dataset).
• 3D Datasets:
  • Het effect was hier nog sterker. Patch-grootte 1 presteerde over het algemeen het beste.
  • Grote patches (14 en 28) leverden aanzienlijk slechtere resultaten op.
  • Verbetering: Er werd een toename in Balanced Accuracy gezien tot 23,78% (bijv. patch 1 vs. 14 op de VesselMNIST3D dataset).
• Computationele Kosten:
  • De verbetering komt met een prijs: het verkleinen van de patch-grootte van 28 naar 1 verhoogt de rekentijd (GFLOPs) drastisch.
  • Bij 2D stijgt de kostenfactor ongeveer 4x per halvering van de patch.
  • Bij 3D is dit effect extremer: het verkleinen van patch 28 naar 1 verhoogt de kosten van ~0,4 GFLOPs naar >800 GFLOPs per volume (een factor van 2000+).
• Attention Maps: Visuele analyse van de attention maps toonde aan dat modellen met kleinere patches zich gerichter richten op diagnostisch relevante gebieden, terwijl modellen met grote patches een uniformer en minder informatief patroon vertonen.

Betekenis en Conclusie

De studie concludeert dat er een duidelijke relatie bestaat tussen patch-grootte en modelprestaties in de medische beeldanalyse. Finere tokenisatie (kleinere patches) stelt ViTs in staat om meer lokale en klinisch relevante kenmerken vast te leggen, wat essentieel is voor complexe medische taken.

Hoewel de rekenkosten voor 3D-data aanzienlijk stijgen, biedt de studie een belangrijke praktische inzichten:

1. Voor medische toepassingen waar nauwkeurigheid cruciaal is, moet de voorkeur worden gegeven aan kleinere patch-groottes (bijv. 1 of 2), zelfs als dit meer rekenkracht vereist.
2. Het gebruik van een ensemble van modellen met verschillende kleine patch-groottes kan de prestaties verder optimaliseren.
3. Het is haalbaar om dergelijke diepgaande analyses uit te voeren zonder toegang te hebben tot enorme clusters van GPU's, zolang de dataset-resolutie maar zorgvuldig wordt geselecteerd.

De auteurs wijzen erop dat toekomstig onderzoek zich moet richten op de toepasbaarheid van deze bevindingen op echte, hoge-resolutie klinische datasets en op het managen van de hoge rekenkosten voor 3D-toepassingen in real-time klinische omgevingen.