Technische Samenvatting: CPS4 – Class Prompt-gestuurde Semi-Supervised Rugsegmentatie met Class-specifieke Consistentiebeperking

Probleemstelling

Automatische rugsegmentatie, specifhankelijk de multi-class segmentatie van wervellichamen (vertebral bodies, VBs) en tussenwervelschijven (intervertebral discs, IVDs) in MRI, is cruciaal voor orthopedische diagnose en chirurgische planning. Hoewel deep learning-methoden aanzienlijk zijn gevorderd, zijn ze afhankelijk van pixel-niveau annotaties, die arbeidsintensief zijn en deskundige kennis vereisen. Semi-supervised learning (SSL) biedt een oplossing door beperkte gelabelde data te gebruiken naast overvloedige ongelabelde data. Bestaande SSL-methoden, die doorgaans vertrouwen op consistentie-regularisatie en pseudo-labeling, lijden echter vaak aan kwalitatief lage pseudo-labels. Deze ruisgevoelige labels introduceren onjuiste supervisie, waardoor modellen zich richten op irrelevante kenmerken en de effectieve exploitatie van ongelabelde data hinderen.

Hoewel Vision-Language Models (VLMs) potentieel tonen voor het verbeteren van de kwaliteit van pseudo-labels door tekstuele class prompts te gebruiken om segmentatiekaarten te genereren, missen huidige VLMs expliciete beperkingen om consistentie tussen class prompts en specifieke objecten in afbeeldingen te waarborgen. Dit gebrek leidt tot een zwakke class-specifieke semantische binding en cross-class interferentie, wat resulteert in een onbevredigende prestatie bij de generatie van multi-class segmentatiekaarten.

Methodologie

De auteurs stellen CPS4 voor, een nieuw framework dat class prompts integreert in semi-supervised rugsegmentatie via een tweestaps trainingsstrategie.

Fase 1: Class-Specifieke Consistentie Beperkte VLM Pretraining

Bestaande VLMs (bijv. CLIP, MedCLIP) voeren cross-modale uitlijning uit op beeldniveau, wat faalt in het waarborgen dat verschillende class prompts zich richten op hun corresponderende objecten. Om dit aan te pakken, introduceert CPS4 een pretraining-fase met behulp van een Vision-Language Model (specifiek gebruikmakend van ViT voor visie en BioClinicalBERT voor tekst) met twee nieuwe verliesfuncties om consistentie tussen class prompts en rugeenheden af te dwingen:

1. Token-Level Attention Loss ($L_{token}$): Deze loss optimaliseert de uitlijning tussen een class prompt token ($F_{T,i}$) en de corresponderende image patch tokens ($F_I$). Het berekent een gelijkenisscore om een grove attention map ($A_i$) te genereren en aggregeert activaties naar de voorspelde rugeenheid regio ($B_i$). Het doel is om de tekstuele prompt nauw te laten koppelen aan de doelrugeenheid in de semantische ruimte.
2. Pixel-Level Attention Loss ($L_{pixel}$): Om te voorkomen dat het model de aandacht te veel aggregeert naar slechts bepaalde subregio's van een doelunit, wordt een pixel-level cross-entropy objectief toegevoegd. Dit zorgt ervoor dat de attention map de gehele doelrugeenheid uniform dekt.

De totale pretraining loss is het gemiddelde van $L_{token}$ en $L_{pixel}$.

Fase 2: Class Prompt Gedreven Semi-Supervised Rugsegmentatie

In deze fase worden de gepretrainde vision-text encoders gebruikt om het pseudo-labeling proces voor ongelabelde data te verbeteren.

• Class Prompt-Gestuurde Segmentatiekaart Generatie: Voor elke ongelabelde rugafbeelding genereert het systeem een binaire segmentatiekaart voor elke class (bijv. "Lumbar Vertebra 1") door de dot product te berekenen tussen image patch tokens en class prompt tokens. Deze kaarten worden geupsampled en omgezet naar binaire masks met behulp van de OTSU automatische drempelmethode. Alle class-specifieke kaarten worden gestapeld om een uniforme multi-class segmentatiekaart ($y^{text}$) te vormen.
• Pseudo Label Generatie & Optimalisatie: Het framework maakt gebruik van een Mean Teacher (MT) architectuur met een student- en een teacher-netwerk. Het teacher-netwerk genereert pseudo-labels ($y^t_u$) voor ongelabelde afbeeldingen via Exponential Moving Average (EMA).
  • Supervised Loss ($L_{sup}$): Berekend tussen de student-predicties op gelabelde data en de ground truth.
  • Semi-Supervised Loss ($L_{semi}$): Bestaat uit twee delen:
    1. $L^{text}_{semi}$: De consistentie loss tussen de student-predictie op ongelabelde data en de class prompt-gegenereerde kaart ($y^{text}$).
    2. $L^{enhance}_{semi}$: De consistentie loss tussen de student-predictie en een verbeterd pseudo-label dat is gecreëerd door de teacher's pseudo-label ($y^t_u$) en de class prompt map ($y^{text}$) te integreren.

Belangrijkste Bijdragen

1. Eerste Tekst-Gestuurde SSL voor Rugsegmentatie: Dit werk is de eerste die tekstuele class prompts introduceert om de kwaliteit van pseudo-labels te verbeteren in semi-supervised rugsegmentatie.
2. Nieuwe Class-Specifieke Consistentiebeperking: De auteurs stellen een specifieke pretraining-mechanisme voor met token- en pixel-level attention losses om expliciete beperkingen af te dwingen, wat zorgt voor een nauwe semantische koppeling tussen elke class prompt en de doelrugeenheid.
3. State-of-the-Art Prestaties: De methode demonstreert superieure prestaties over diverse gelabelde data-ratio's, met name in scenario's met weinig data.

Experimentele Resultaten

De methode werd geëvalueerd op de MRSpineSeg dataset, die T2-gewogen MR volumetrische beelden bevat van 215 subjecten met 20 categorieën (19 rugstructuren en achtergrond).

• Prestaties: Gebruikmakend van slechts 5% gelabelde data, bereikte CPS4 een Dice score van 80,44% en een mIoU van 70,63%.
• Vergelijking: Deze resultaten overtreffen populaire semi-supervised learning methoden (bijv. MT, BCP, SS-Net) en bestaande VLM-gebaseerde benaderingen (bijv. CLIP, MedCLIP, MGCA, EGMA).
  • Vergeleken met de suboptimale GraphCL methode met tekst prompts op 5% gelabelde data, verbeterde CPS4 de mDice met 2,23% en de mIoU met 2,42%.
• Robuustheid: Ablatie-studies bevestigden dat zowel de token-level als de pixel-level losses significant bijdragen aan de prestaties. De methode vertoonde stabiele prestatiewinst naarmate de ratio van gelabelde data toenam van 5% naar 50%, waarbij de MT baseline met een duidelijke marge werd verslagen.
• Visualisatie: De door CPS4 gegenereerde attention maps vertoonden een sterke uitlijning met de doelrugunits, terwijl baseline VLMs (zoals MedCLIP) verspreide activaties en een zwakke semantische binding vertoonden.

Betekenis

Het artikel claimt dat CPS4 effectief het probleem van labeltekort bij rugsegmentatie mitigeert door de semantische kracht van tekst te benutten om de generatie van hoogwaardige pseudo-labels te sturen. Door tijdens de pretraining de consistentie tussen class prompts en rugeenheden expliciet te beperken, overwint het framework de beperkingen van huidige VLMs in medische beeldvormingstaken. De resultaten suggereren dat deze aanpak een meer betrouwbare class prompt-gestuurde pseudo-label generatie mogelijk maakt, wat de effectiviteit van semi-supervised learning in scenario's met beperkte geannoteerde data aanzienlijk verbetert. De auteurs merken op dat hun code beschikbaar zal worden gesteld om verder onderzoek te faciliteren.