Technische Samenvatting: Locus – Anatomie-gestuurde Attention Regularisatie voor Medische Beeldclassificatie

1. Probleemstelling

Medische beeldclassificatiemodellen worden doorgaans getraind met enkel beeldniveau-labels, zonder expliciete ruimtelijke supervisie. Als gevolg hiervan leren deze modellen vaak te vertrouwen op statistische correlaties met anatomisch irrelevante regio's, zoals beeldartefacten, omliggend weefsel of dataset-specifieke contextuele aanwijzingen, in plaats van op de diagnostisch relevante anatomische structuren. Dit gedrag beperkt de interpreteerbaarheid en generaliseerbaarheid van het model.

Hoewel eerdere benaderingen hebben geprobeerd dit aan te pakken door anatomische vorminformatie (bijv. segmentatiemaskers) te integreren, kampen zij met significante beperkingen:

• Handmatige Annotatie: Het vertrouwen op door experts getekende maskers is kostbaar en moeilijk schaalbaar.
• Gezamenlijke Training (Joint Training): Het trainen van een segmentatie-tak naast een classifier introduceert optimalisatieproblemen, zoals concurrerende gradiënten en ruizige voorspellingen die de sturing kunnen corrumperen.
• CAM-gebaseerde Pseudo-maskers: Hoewel deze methode de annotatiekosten vermijdt, activeren ze vaak alleen de discriminatieve regio's die gecorreleerd zijn met het label, in plaats van de volledige relevante anatomische structuur.
• Dedicated Segmentatiemodellen: Het trainen van een nieuw segmentatiemodel vanaf nul introduceert opnieuw de last van annotatie en computationele overhead.

2. Methodologie: Het Locus-framework

De auteurs stellen Locus voor, een framework dat gebruikmaakt van vooraf getrainde Segmentation Foundation Models (SFMs) (bijv. Medical SAM, BiomedParse) om de aandacht van de classifier te sturen zonder dat er specifieke segmentatietraining of handmatige annotaties nodig zijn.

2.1 Anatomie-centrische Attention Regularisatie

In plaats van een rigide pixel-gewijze uitlijning tussen de attention map van de classifier en het anatomische masker af te dwingen (wat te restrictief kan zijn), introduceert Locus een hinge-gebaseerde regularisatieterm ($R$).

• Mechanisme: Het framework berekent het gewogen gemiddelde van de attention-energie binnen de anatomische voorgrond ($A_{fg}$) en de achtergrond ($A_{bg}$) afgeleid van het SFM-masker.
• Verliesfunctie: De regularisator wordt gedefinieerd als $R(M, A) = \max(0, A_{bg} - A_{fg})$.
• Logica: Deze term straft de classifier alleen af wanneer de achtergrond-attention domineert boven de voorgrond-attention. Als de voorgrond-attention al voldoende is ($A_{fg} \geq A_{bg}$), blijft de regularisator inactief. Deze aanpak voorkomt het overmatig onderdrukken van potentieel diagnostisch relevante contextuele informatie, terwijl het ervoor zorgt dat het model niet primair leunt op achtergrondruis.

2.2 Patch-wise Anatomische Prompting

Standaard prompting van SFMs op volledige beelden faalt vaak in het vastleggen van alle relevante anatomische structuren, vooral wanneer deze klein of verspreid zijn (bijv. verspreide laesies of bilaterale organen), waarbij vaak alleen de meest dominante regio wordt teruggegeven.

• Strategie: Locus verdeelt het invoerbeeld in een raster van niet-overlappende patches. Elke patch wordt onafhankelijk geprompt met dezelfde klasse-afgeleide tekstbeschrijving (bijv. "long", "tumor").
• Aggregatie: De resulterende patch-niveau maskers worden geaggregeerd via een unie om een anatomisch masker met volledige resolutie ($M$) te vormen.
• Doel: Deze strategie maximaliseert de anatomische dekking om ervoor te zorgen dat alle relevante regio's bijdragen aan de attention-supervisie, zelfs als het resulterende masker grof is. De auteurs benadrukken dat precieze pixel-niveau segmentatie niet vereist is; het doel is voldoende anatomische verankering om de aandacht te sturen.

2.3 Netwerkverlies

Het totale netwerkverlies combineert de standaard taakverlies (cross-entropy) met de anatomie-centrische regularisator:
$$L_{net} = L_{task} + \lambda \cdot R(M, A)$$

• Trainingsstrategie: Er wordt een warm-up strategie gebruikt waarbij het netwerk gedurende enkele epochs eerst alleen met $L_{task}$ wordt getraind, waarna de sterkte van de regularisatie ($\lambda$) geleidelijk wordt opgevoerd.
• Flexibiliteit: Het framework is agnostisch voor de specifieke methode van extractie van de attention (bijv. Grad-CAM, Grad-CAM++, Vanilla CAM) en de keuze van de SFM.

3. Belangrijkste Bijdragen

1. Nieuw Regularisatieframework: Stelt een anatomie-centrische attention regularisator voor die de aandacht adaptief balanceert tussen voorgrond en achtergrond, waarbij de classifier alleen wordt gestraft wanneer de achtergrond-attention domineert, waardoor de noodzaak voor dedicated segmentatietraining wordt weggenomen.
2. Plug-and-Play Aanpasbaarheid: Het framework is eenvoudig aan te passen aan diverse classifier backbones (CNN's en Transformers) en SFM's zonder de architectuur van de classifier te wijzigen of de foundation model te finetunen.
3. Uitgebreide Validatie: Gevalideerd op acht diverse medische beeldvormingsdatasets (dermoscopie, röntgen, histopathologie, echografie en cardiale MRI), waarbij consistente prestatiewinsten en verbeterde interpreteerbaarheid werden aangetoond.

4. Experimentele Resultaten

De auteurs evalueerden Locus tegenover drie baselines: standaard Cross-Entropy (CE), Masked input (gebruikmakend van SFM-maskers om de achtergrond te verwijderen) en MGA (Dice-gebaseerde attention alignment die expert-annotaties vereist).

• Classificatieprestaties: Locus behaalde de hoogste nauwkeurigheid en F1-scores in bijna alle settings over zes benchmark-datasets (HAM10000, PneumoniaMNIST, BTXRD, PanNuke, ChestX-ray8) en een dataset met cardiale cine MRI. Het presteerde consequent beter dan onbeperkte training en de masked-input baselines.
• Kwaliteit van de Attention: Visualisaties (Grad-CAM) toonden aan dat Locus meer gefocuste en anatomisch coherente attention maps produceerde. Opvallend genoeg, voor architecturen die gevoelig zijn voor achtergronddominantie (bijv. ConvNeXt, ViT), maakte Locus de attention aanzienlijk scherper op relevante structuren (bijv. longvelden, laesies) vergeleken met standaard CE-training.
• Robuustheid: Onder testtijd-perturbaties (helderheid, contrast, Gaussische vervaging, ruis) behielden modellen getraind met Locus een hogere nauwkeurigheid dan CE-only baselines, wat suggereert dat de regularisatie het vertrouwen in spurieuze kenmerken die gevoelig zijn voor verstoringen ontmoedigt.
• Generaliseerbaarheid: De voordelen van de regularisator waren consistent over verschillende methoden voor extractie van de attention (Grad-CAM, Grad-CAM++, Vanilla CAM, Gradient Norm) en verschillende netwerk backbones (ResNet, ConvNeXt, ViT).

5. Betekenis en Claims

Het artikel claimt dat Locus een praktische oplossing biedt voor het "black box"-karakter van medische beeldclassificatiemodellen door ze te sturen naar diagnostisch relevante anatomische structuren zonder de hoge kosten van handmatige annotatie of de complexiteit van gezamenlijke segmentatie- en classificatietraining.

Door gebruik te maken van de "out-of-the-box" lokalisatiecapaciteiten van foundation modellen en een flexibele, hinge-gebaseerde regularisator toe te passen, bereikt het framework:

• Verbeterde Interpreteerbaarheid: Modellen focussen op de juiste anatomische regio's (bijv. myocard in cardiale MRI) in plaats van op artefacten of de achtergrond.
• Verhoogde Generaliseerbaarheid: Betere prestaties over diverse modaliteiten en robuustheid tegen beeldperturbaties.
• Efficiëntie: Het vermijdt de computationele overhead van het trainen van een dedicated segmentatiemodel terwijl het betrouwbare anatomische sturing biedt.

De auteurs concluderen dat deze aanpak een belangrijke stap is naar het betrouwbaarder en klinisch interpreteerbaar maken van medische AI, met name in scenario's waar pixel-niveau annotaties niet beschikbaar zijn. Toekomstig werk wordt gesuggereerd om onzekerheid-bewuste weging en end-to-end leren van patch-configuraties te verkennen.