Unsupervised Source-Free Ranking of Biomedical Segmentation Models Under Distribution Shift

Probleemstelling

In de biomedische beeldanalyse is segmentatie cruciaal voor het interpreteren van biologische structuren. Hoewel diep leren de nauwkeurigheid heeft verbeterd, blijft de noodzaak voor handmatige annotatie van trainingsdata een grote bottleneck. Een veelbelovende oplossing is het hergebruiken van vooraf getrainde modellen (bijvoorbeeld uit repositories zoals de BioImage Model Zoo) voor nieuwe doel-datasets.

Echter, het selecteren van het meest geschikte model voor een nieuwe dataset is moeilijk, vooral omdat:

1. Geen labels beschikbaar zijn: In de praktijk ontbreekt vaak ground-truth data in het doel-domein om modellen te evalueren.
2. Bestaande methoden falen: Veel bestaande methoden voor het schatten van "transferability" (overdraagbaarheid) zijn ontworpen voor classificatie en vereisen labels, toegang tot de feature-space, of specifieke trainingsaannames.
3. Instance vs. Semantic Segmentatie: Bestaande methoden werken vaak niet voor instance segmentation, omdat de output (een verzameling van objecten) geen vaste correspondentie heeft met de feature-space, in tegenstelling tot pixel-classificatie.
4. Domeinverschuiving (Distribution Shift): Modellen getraind op bron-domeinen presteren vaak slecht op doel-domeinen met andere beeldkwaliteit of anatomische variaties.

Er is dus behoefte aan een methode die onzelftoezicht (unsupervised), bron-vrij (source-free) en model-agnostisch is, en die werkt voor zowel semantische als instance segmentatie zonder dat labels of toegang tot de interne architectuur nodig zijn.

Methodologie: Consistency-Based Model Ranking (CMR)

De auteurs introduceren CMR, een raamwerk dat modellen rangschikt op basis van de consistentie van hun voorspellingen onder perturbaties (verstoringen). Het kernidee is dat een goed getransfereerd model robuuste beslissingsgrenzen heeft; kleine veranderingen in de invoer of features leiden niet tot grote veranderingen in de output.

Kerncomponenten:

1. Perturbatie-strategie:

   • Input Space: Toepassing van test-time augmentaties (TTA) zoals Gaussisch ruis, gamma-correctie, helderheids- en contrastaanpassingen.
   • Feature Space: Toepassing van Test-Time Dropout (TTD) op de features van het model tijdens inferentie (zonder het model opnieuw te trainen). Dit kan gebeuren op de bottleneck, skip-connections, of alle lagen.

2. Consistentie-maatstaven:
   De methode vergelijkt de voorspelling van het oorspronkelijke beeld ($\hat{y}$) met die van het verstoord beeld ($\tilde{y}$).

   • Voor Semantische Segmentatie:
     • CMR-EI (Soft): Gebaseerd op Effective Invariance. Meet de overeenkomst in softmax-kansen en labels, gewogen door de voorgrond-pixels.
     • CMR-NHD (Hard): Gebaseerd op de genormaliseerde Hamming-afstand. Telt het aantal pixel-label flips in de voorgrond. Dit is equivalent aan een IoU-maatstaf maar gericht op stabiliteit.
   • Voor Instance Segmentatie:
     • Omdat instance-labels willekeurig zijn en geen vaste volgorde hebben, kunnen pixels niet direct worden vergeleken.
     • De auteurs gebruiken de Adapted Rand Score (ARS). Deze meet de overeenkomst tussen paren van pixels: als twee pixels in de verstoord voorspelling tot hetzelfde object behoren, moeten ze dat ook doen in de oorspronkelijke voorspelling (en vice versa). Dit omzeilt het probleem van willekeurige instance-ID's.

3. Rangschikking:
   Voor elk model wordt een consistentiescore berekend over meerdere perturbaties. De modellen worden vervolgens aflopend gerangschikt op deze score. Een hoge consistentie impliceert een betere prestatie op het doel-domein.

Belangrijkste Bijdragen

1. Eerste Unsupervised Methode: Het is de eerste methode die zowel semantische als instance segmentatiemodellen kan rangschikken zonder labels, bron-data of toegang tot de interne features van het model (black-box compatibel).
2. Model-agnostisch: Werkt met elke architectuur (U-Net, Transformers, SAM-varianten) en is geschikt voor repositories met heterogene modellen.
3. Uitgebreide Validatie: De methode is getest op een breed scala aan biomedische taken (kernen, cellen, mitochondriën) in zowel 2D als 3D, inclusief lichtmicroscopie en elektronenmicroscopie.
4. Toepassing op UDA: De methode werkt ook na Unsupervised Domain Adaptation (UDA), wat essentieel is om te bepalen of een adaptatiemethode het model daadwerkelijk heeft verbeterd.
5. Benchmark Resultaten: De methode slaagt erin om de rangschikking van het ToothFairy2-challenge (3D tandheelkundige segmentatie) te reproduceren zonder toegang tot de testlabels.

Resultaten

De auteurs hebben CMR geëvalueerd tegenover bestaande state-of-the-art methoden (zoals Transfer Score, LEEP, LogME, Dispersion) op diverse datasets.

• Semantische Segmentatie: CMR toonde een sterke correlatie (Pearson, Spearman, Kendall tau) met de werkelijke prestaties (F1-score) op doel-datasets. Bestaande methoden, zelfs die met labels, faalden vaak bij biomedische data of bij de overstap van classificatie naar segmentatie.
  • Voorbeeld: Op mitochondriën-datasets bereikte CMR een Kendall tau van ~0.77-0.87, terwijl Transfer Score (TS) vaak rond de 0.25 bleef.
• Instance Segmentatie: CMR-ARS presteerde aanzienlijk beter dan de enige andere bestaande methode voor instance ranking (SEG), die gebaseerd is op objectcentroïden en minder gevoelig is voor pixel-kwaliteit. CMR-ARS behaalde correlaties van ~0.69-0.83, terwijl SEG vaak wisselvallig was.
• Robuustheid: De methode is stabiel over een breed scala aan perturbatiesterktes (ruis, dropout).
• Efficiëntie: CMR vereist slechts twee inferentie-passes per afbeelding (oorspronkelijk + één perturbatie), wat veel efficiënter is dan ensemble-methoden of methoden die Monte-Carlo dropout vereisen (vaak 10+ passes).
• 3D & Multi-class: Op de complexe 3D ToothFairy2-dataset (42 klassen) behaalde CMR correlaties tot 1.0, wat aantoont dat het werkt voor zware, multi-class taken.

Significantie en Conclusie

Dit paper biedt een praktische oplossing voor een groot probleem in de biomedische AI: hoe kies je het beste model uit een repository zonder dure annotatie?

• Praktische Toepasbaarheid: Het stelt onderzoekers en clinici in staat om zelf hun modellen te valideren en te selecteren voordat ze worden ingezet, zelfs als ze geen labels hebben.
• Universeel: Door het werken op output-niveau (in plaats van feature-niveau) is de methode toepasbaar op elke moderne architectuur, inclusief de nieuwste "generalist" modellen zoals µSAM en Cellpose-SAM.
• Toekomstperspectief: De methode faciliteert schaalbaar hergebruik van modellen en benchmarking in groeiende model-repositories, wat de adoptie van deep learning in domeinen met schaarse data (zoals de biomedische wetenschap) zal versnellen.

Kortom, CMR is een krachtig, efficiënt en betrouwbaar instrument voor het onafhankelijk beoordelen van de kwaliteit van segmentatiemodellen onder domeinverschuiving.