Scaling In-Context Segmentation with Hierarchical Supervision

Hieronder volgt een gedetailleerde technische samenvatting van het paper "Scaling In-Context Segmentation with Hierarchical Supervision" in het Nederlands.

1. Het Probleem

In-Context Learning (ICL) is een veelbelovende aanpak voor medische beeldsegmentatie, waarbij modellen zich kunnen aanpassen aan nieuwe anatomische structuren op basis van slechts een paar voorbeelden, zonder dat ze opnieuw getraind hoeven te worden. Dit vermindert de last van klinische annotatie aanzienlijk.

Echter, het toepassen van bestaande ICL-methoden op hoogresolutie medische volumes stuit op twee grote problemen:

• Rekenkundige schaalbaarheid: Standaard methoden (zoals UniverSeg) maken gebruik van dichte, globale cross-attention. De rekenkosten hiervan groeien kwadratisch met de beeldresolutie ($O((HW)^2)$), wat ze onpraktisch maakt voor hoge resoluties (bijv. 512x512 of groter).
• Gebrek aan expliciete supervisie: Bestaande methoden die proberen selectief te zijn (zoals Tyche of EICSeg), leren de selectie van relevante gebieden impliciet via de segmentatieverliezen alleen. Dit leidt vaak tot redundante berekeningen in niet-informatieve achtergrondgebieden, omdat het model niet expliciet wordt gestuurd om irrelevante delen van het beeld te negeren voordat zware verwerking plaatsvindt.

2. Methodologie: PatchICL

De auteurs stellen PatchICL voor, een hiërarchisch raamwerk dat selectieve patch-sampling combineert met multi-level supervisie. Het doel is om actief alleen de meest informatieve anatomische regio's te identificeren en te verwerken.

Kerncomponenten:

• Coarse-to-Fine Cascade: Het model verwerkt afbeeldingen via een cascade van $M$ resolutieniveaus (bijv. van 24x24 naar 128x128).
• Entropie-gestuurde Patch Sampling (Doelbeeld):
  • Op elk niveau $\ell > 1$ worden patches geselecteerd op basis van de onzekerheid (entropie) van de voorspelling van het vorige niveau.
  • De entropie $H(p)$ wordt berekend over de voorspelde waarschijnlijkheid. Gebieden met hoge entropie (onzekere regio's) krijgen een hogere weging.
  • Er wordt gebruikgemaakt van Gumbel-top-K sampling om een discrete selectie van $K$ patches te maken die het meest informatief zijn.
• Context Sampling: Voor context-afbeeldingen (de voorbeelden) worden patches geselecteerd op basis van de nabijheid tot de randen van de grondwaarheid (ground truth). Dit zorgt ervoor dat de context patches de meest structureel informatieve gebieden dekken.
• Multi-Level Supervisie: In tegenstelling tot impliciete methoden, krijgt het selectieproces expliciete supervisie. Elke resolutieniveau heeft zijn eigen verliesfunctie (combinatie van Binary Cross-Entropy en Dice Loss) die wordt toegepast op de gecombineerde voorspelling. Dit dwingt het model om al vroeg in de cascade onnuttige achtergrond te verwerpen.
• Architectuur:
  • Een bevroren encoder (van UniverSeg) extrahert features.
  • Een Transformer-laag verwerkt doel- en context-patches gezamenlijk met 2D Rotary Position Embeddings (RoPE) en bidirectionele attention.
  • Een CNN-decoder reconstructeert de segmentatie voor elke patch.
  • Voorspellingen van overlappende patches worden gemiddeld en via additieve fusie gecombineerd met de opgezoomde voorspelling van het vorige niveau.

3. Belangrijkste Bijdragen

1. Hiërarchisch Raamwerk: Introductie van een coarse-to-fine cascade met entropie-gestuurde Gumbel-top-K patch-sampling, waardoor berekeningen zich focussen op onzekere regio's.
2. Expliciete Supervisie: Toepassing van multi-level supervisie op elk resolutieniveau, wat zorgt voor een direct trainingsignaal voor het selectieproces en het vroegtijdig verwijderen van irrelevante achtergrond mogelijk maakt.
3. Efficiëntie en Schaalbaarheid: Demonstratie dat PatchICL concurrerende prestaties behaalt op CT-segmentatie binnen het domein, terwijl de rekenkosten met 44% worden verlaagd bij een resolutie van 512x512 in vergelijking met de sterke baseline UniverSeg.

4. Resultaten

De methoden zijn geëvalueerd op drie generalisatie-scenario's:

• In-Domain (CT):
  • Op de TotalSegmentator CT-dataset (met 45 vastgehouden anatomische klassen) behaalt PatchICL een hogere gemiddelde Dice-score (0.417) dan UniverSeg (0.397).
  • Bij een resolutie van 512x512 verbruikt PatchICL slechts 308 GFLOPs per sample, tegenover 549 GFLOPs voor UniverSeg (een besparing van 44%), terwijl de nauwkeurigheid behouden blijft of verbetert.
• Cross-Modality (MRI):
  • Op de TotalSegmentator MRI-dataset (getraind op CT, getest op MRI) presteert UniverSeg iets beter overall (+2.9% Dice), vooral bij complexe articulaties zoals de wervelkolom. PatchICL presteert echter nog steeds competitief.
• Out-of-Domain (MedSegBench):
  • Op 35 diverse datasets (ultrasound, X-ray, OCT, dermoscopy, etc.) presteert PatchICL beter dan UniverSeg op 6 van de 13 modaliteiten.
  • De grootste winst is zichtbaar bij modaliteiten met gelokaliseerde pathologie en fijne texturen, zoals OCT (+13.2% Dice) en Dermoscopy (+6.6% Dice).
  • UniverSeg doet het beter bij modaliteiten met hoog contrast en globaal gestructureerde doelen (zoals X-Ray en Nuclear Cell).

5. Betekenis en Conclusie

PatchICL lost het fundamentele probleem van schaalbaarheid in ICL voor medische beeldvorming op door de kwadratische complexiteit van dichte attention te doorbreken. Door expliciet te leren welke beelddelen relevant zijn via multi-level supervisie, kan het model rekenkracht besparen zonder in te leveren op de nauwkeurigheid.

De studie toont aan dat voor hoogresolutie taken en modaliteiten met gelokaliseerde pathologie (zoals oog- of huidafbeeldingen), een selectieve, hiërarchische aanpak superieur is aan globale attention-mechanismen. Dit maakt ICL praktischer voor klinische toepassingen waar hoge resoluties en snelle verwerking essentieel zijn, terwijl de last van annotatie wordt verlaagd. De auteurs merken op dat toekomstig werk zich zal richten op het leren van betere sampling-strategieën en adaptieve niveau-selectie.