Domain Adaptation Without the Compute Burden for Efficient Whole Slide Image Analysis

Probleemstelling

De analyse van Whole Slide Images (WSI's) in de computergestuurde pathologie is cruciaal voor vroege diagnose en behandeling van tumoren. WSIs hebben echter een extreme resolutie (gigapixels), waardoor end-to-end training van diepe leermodellen onpraktisch is vanwege geheugenbeperkingen.

• Huidige aanpak: De standaardmethode is het gebruik van Multiple Instance Learning (MIL) met vaste (frozen) patch-representaties. Deze worden vooraf berekend met beeldencoders die vaak zijn voorgetraind op natuurlijke afbeeldingen (bijv. ImageNet).
• Beperkingen:
  1. Domain Gap: ImageNet-encoders missen domeinspecifieke kenmerken van histopathologie.
  2. Rekenkundige last: Domeinspecifieke voortraining op histopathologische data is uiterst duur en vereist zware GPU-resources.
  3. Gebrek aan aanpassing: Zelfs met domeinspecifieke encoders kunnen modellen geen taakspecifieke kenmerken leren uit de ruwe data omdat de patch-representaties vast staan.
  4. Sampling-probleem: End-to-end training vereist vaak het verwerken van duizenden patches per slide, wat rekenkundig onhaalbaar is. Bestaande MIL-methoden presteren vaak slecht bij lage sampling-rates (weinig patches).

Methodologie: EfficientWSI (eWSI)

De auteurs stellen eWSI voor, een framework dat Parameter-Efficient Fine-Tuning (PEFT) combineert met een aangepaste MIL-architectuur om end-to-end training mogelijk te maken op algemene GPU's zonder zware voortraining.

Het framework bestaat uit drie kerncomponenten:

1. Patch Sampling:
   In plaats van alle patches te verwerken, wordt er tijdens elke training iteratie een willekeurige, spaarzame subset van $M$ patches uit de WSI gehaald (bijv. 64 tot 512 patches). Dit verlaagt de rekenlast drastisch.

2. PEFT Adaptatie (LoRA):
   De beeldencoder (bijv. ViT) wordt niet volledig opnieuw getraind. In plaats daarvan wordt Low-Rank Adaptation (LoRA) gebruikt.

   • De gewichten van de encoder worden bevroren.
   • Alleen specifieke lagen worden aangepast: de query-projecties ($q$) en de feed-forward lagen (MLP).
   • Dit stelt het model in staat om domeinspecifieke informatie te leren van de ruwe data, maar met een zeer klein aantal trainbare parameters.

3. MIL Aggregatie (LinMax):
   Traditionele attention-based MIL-methoden (zoals ABMIL) presteren slecht bij lage sampling-rates. De auteurs introduceren LinMax, een nieuwe aggregator:

   • Bestaat uit een stapel van $L$ volledig verbonden lagen met $tanh$-activaties.
   • Gevolgd door een max-pooling operatie over de patch-dimensie.
   • Dit ontwerp behoudt de robuustheid van max-pooling (in lijn met de MIL-aanname dat ten minste één positieve patch een positieve slide betekent) maar verhoogt de representatieve capaciteit door de diepte van de lagen, zonder de parametercount te verhogen.

Belangrijkste Bijdragen

• Efficiënt End-to-End Training: eWSI maakt het mogelijk om WSIs end-to-end te trainen op een enkele algemene GPU (bijv. RTX 2080Ti/3090), zonder de noodzaak van dure domeinspecifieke voortraining.
• Robuuste MIL voor Lage Sampling: De introductie van LinMax lost het probleem op dat bestaande MIL-methoden instabiel worden bij het gebruik van weinig patches, wat essentieel is voor de efficiëntie van eWSI.
• Domeinadaptatie zonder Overhead: Het framework demonstreert dat het combineren van PEFT met ImageNet-encoders prestaties kan leveren die vergelijkbaar zijn met of beter zijn dan methoden met zware domeinspecifieke voortraining.
• Open Source: De auteurs publiceren hun code, voorgetrainde modellen (inclusief een eigen Hist-SSL encoder) en checkpoints.

Resultaten

De auteurs evalueerden eWSI op zeven taken over drie datasets: Camelyon16 (tumor metastase), TCGA (kankersubtypes en moleculaire eigenschappen) en BRACS (goedaardig/atypisch/kwaadaardig).

• Prestaties:
  • eWSI met ImageNet-encoders (iNet-Sup) presteert aanzienlijk beter dan standaard MIL-methoden met vaste encoders en bereikt vergelijkbare resultaten met methoden die gebruikmaken van domeinspecifieke voortraining.
  • Wanneer eWSI wordt toegepast op domeinspecifieke encoders (Hist-SSL), wordt de prestatie in de meeste gevallen verder verbeterd, wat aantoont dat het model nuttige taakspecifieke informatie kan leren.
  • Op Camelyon16 bereikte eWSI met 512 patches en een ImageNet-encoder een AUC van 93.4%, wat beter is dan de beste bestaande MIL-methoden (bijv. ACMIL op 84.1% met dezelfde encoder).
• Robuustheid: Het model blijft robuust presteren zelfs bij zeer lage sampling-rates (bijv. 64 patches), terwijl attention-based MIL-methoden hier sterk op achteruitgaan.
• Berekening:
  • eWSI is aanzienlijk sneller dan bestaande end-to-end methoden (zoals StreamingCLAM of C2C). Training van 100 epoch op Camelyon16 duurt ongeveer 2,5 uur (bij M=512) op een RTX 3090, vergeleken met 45+ uur voor andere end-to-end methoden.
  • Hoewel het iets langer duurt dan het fine-tunen van een vaste encoder, is de kosten-batenverhouding (prestatie vs. tijd) zeer gunstig.

Betekenis en Conclusie

Het paper toont aan dat de hoge rekenkundige barrière voor end-to-end WSI-analyse kan worden doorbroken door slimme integratie van PEFT en een aangepaste aggregatiestrategie.

• Praktische Impact: Pathologie-instituten kunnen nu taakspecifieke modellen trainen op hun eigen hardware zonder afhankelijk te zijn van grote foundation models of zware voortrainingspipelines.
• Wetenschappelijke Inzichten: Het paper benadrukt dat voor WSI-taken de keuze van de MIL-architectuur (max-pooling vs. attention) cruciaal is bij lage sampling-rates en dat PEFT effectief kan worden gebruikt voor domeinadaptatie in de histopathologie, zelfs vanuit ImageNet-voorbereiding.

Kortom, eWSI biedt een schaalbare, kostenefficiënte en prestatie-gerichte oplossing voor de toekomst van computergestuurde pathologie.