K-MaT: Knowledge-Anchored Manifold Transport for Cross-Modal Prompt Learning in Medical Imaging

K-MaT is een innovatief prompt-learning framework dat de prestaties van grote biomedische visueel-taalmodellen op lage-resolutie medische beeldvorming verbetert door klinische kennis te gebruiken en optimale transporttechnieken toe te passen om modality-specifieke kortsluitingen te voorkomen zonder extra trainingsdata nodig te hebben.

De kern

Stel je voor dat je een zeer ervaren radioloog bent die gespecialiseerd is in het bekijken van 3D-scanbeelden (zoals CT-scan). Hij kan elke ziekte met zijn ogen dicht diagnosticeren omdat hij de subtiele details in die complexe beelden kent.

Nu willen we deze kennis gebruiken om een arts te helpen die werkt met simpele 2D-röntgenfoto's of echografie. Het probleem? De "taal" van de beelden is totaal anders. Een CT-scan ziet eruit als een gedetailleerde 3D-kaas, terwijl een röntgenfoto eruitziet als een platte schaduw.

Als je de CT-expert direct laat kijken naar de röntgenfoto's, raakt hij in de war. Hij probeert patronen te zoeken die alleen in de 3D-kaas bestaan, en hij vergeet zijn algemene medische kennis. Dit noemen onderzoekers "catastrophic forgetting" (catastrofale vergetelheid): de AI wordt zo goed in het lezen van de CT-scan dat ze vergeet hoe ze een röntgenfoto moet lezen.

K-MaT is de slimme oplossing die de auteurs van dit paper hebben bedacht om dit probleem op te lossen. Hier is hoe het werkt, vertaald in alledaagse termen:

1. De "Taal" van de AI (Prompt Learning)

Stel je voor dat de AI een vertaler is. Om een ziekte te herkennen, moet de vertaler een "prompt" (een soort instructiezin) gebruiken.

• Het oude probleem: De AI leerde een instructiezin die perfect paste bij de CT-scan, maar die zin was onbegrijpelijk voor de röntgenfoto.
• De K-MaT oplossing: Ze maken de instructiezin op in twee delen:
  1. Een algemeen medisch deel (wat is een tumor eigenlijk?).
  2. Een specifiek beeld-deel (hoe ziet die tumor eruit op een CT-scan vs. een röntgenfoto?).
     Hierdoor kan de AI de algemene kennis behouden, terwijl ze zich aanpast aan het specifieke type foto.

2. De "Anker" (Knowledge Anchoring)

Stel je voor dat de AI een bootje is dat op een woelige zee drijft (de verschillende beeldtypes). Zonder anker drijft het weg naar een eiland waar alleen CT-scans bestaan.

• De Anker: De auteurs gebruiken een LLM (een slimme taalcomputer) om beschrijvingen van ziekten te schrijven (bijv. "een kwaadaardige massa met onregelmatige randen").
• Deze tekstuele beschrijvingen fungeren als een vast anker in de zee. De AI mag wel aanpassen hoe ze naar de foto kijkt, maar ze mag nooit vergeten wat de tekstuele beschrijving van de ziekte is. Dit zorgt ervoor dat ze niet "dwaalt" en haar medische kennis behoudt.

3. De "Magische Transport" (Manifold Transport)

Dit is het meest creatieve deel. Stel je voor dat je twee verschillende landen hebt:

• Land A (CT-scan): Hier wonen de ziektes in een complexe, 3D-stad.
• Land B (Röntgenfoto): Hier wonen dezelfde ziektes, maar in een platte, 2D-stad.

De AI moet de wegen van Land A naar Land B overbrengen, zonder dat ze de wegen zelf ziet in Land B (want ze heeft geen trainingsdata van de röntgenfoto's!).

• De Oplossing (FGW Optimal Transport): K-MaT gebruikt een wiskundige "magische kaart". Deze kaart kijkt niet naar de straten zelf, maar naar de relaties tussen de gebouwen.
  • Voorbeeld: In Land A is het ziekenhuis altijd rechts van het park. In Land B moet het ziekenhuis ook rechts van het park liggen, ook al zien de gebouwen er anders uit.
  • K-MaT zorgt ervoor dat de "kaart" van de röntgenfoto (Land B) exact dezelfde structuur heeft als de kaart van de CT-scan (Land A). Hierdoor "leert" de AI hoe een ziekte eruitziet op een röntgenfoto, puur door de structuur van de CT-scan te kopiëren, zonder ooit een röntgenfoto te hebben gezien tijdens het leren.

Waarom is dit belangrijk?

In de medische wereld hebben we vaak veel data van dure, geavanceerde apparatuur (zoals MRI of CT), maar weinig data van goedkope, toegankelijke apparatuur (zoals röntgen of echo) in ontwikkelingslanden of huisartsenpraktijken.

• Zonder K-MaT: De AI werkt perfect in het ziekenhuis met de dure scanner, maar faalt volledig bij de huisarts.
• Met K-MaT: De AI kan haar kennis van de dure scanner "transporteren" naar de simpele apparatuur. Ze vergeet niet hoe ze moet diagnosticeren en kan zelfs ziektes vinden op simpele foto's die ze nooit heeft getraind om te zien.

Kort samengevat:
K-MaT is als een slimme vertaler die een boek van een dure, ingewikkelde taal (CT-scan) naar een simpele, dagelijkse taal (röntgenfoto) vertaalt. Hij gebruikt een vast woordenboek (de tekst-analyses) om zeker te weten dat de betekenis niet verloren gaat, en een magische landkaart om de structuur van de complexe taal over te brengen naar de simpele taal, zodat de lezer (de AI) alles begrijpt, zelfs zonder het simpele boek ooit te hebben gelezen.

Technische samenvatting

Probleemstelling

Diepe leermodellen voor medische beeldvorming lijden vaak onder prestatieverlies bij distributieveranderingen, met name bij cross-modale overdracht. Modellen die zijn getraind op hoogwaardige beeldvormingsmodaliteiten (zoals MRI of CT-scans) generaliseren vaak niet betrouwbaar naar toegankelijkere, "laagwaardige" modaliteiten (zoals röntgenfoto's of echografie), ondanks dat ze dezelfde onderliggende pathologieën bestrijken.

Bestaande Vision-Language Models (VLMs) en prompt-learning methoden (zoals CoOp en BiomedCoOp) proberen dit op te lossen door aanpasbare prompts te gebruiken. Echter, in een strikt zero-shot regime (waarbij geen trainingsdata van de doel-modaliteit beschikbaar is), lijden deze methoden aan catastrofaal vergeten. De leerbare prompts "kollapsen" naar modality-specifieke statistieken (shortcuts) van de bron-modaliteit en verliezen de essentiële, gedeelde diagnostische semantiek die nodig is om de doel-modaliteit correct te interpreteren.

Methodologie: K-MaT

De auteurs stellen K-MaT (Knowledge-Anchored Manifold Transport) voor, een prompt-learning framework dat beslissingsstructuren overbrengt van hoogwaardige naar laagwaardige modaliteiten zonder dat er trainingsafbeeldingen van de laagwaardige modaliteit nodig zijn. Het framework bouwt voort op de BiomedCLIP-backbone (met bevroren encoders) en introduceert drie kernmechanismen:

1. Gefactoriseerde Prompts:
   In plaats van één uniforme prompt, worden de prompts gefactoriseerd in:

   • Class-Specific Context (CSC): Specifiek voor de ziekteklasse.
   • Modality-Specific Context (MSC): Specifiek voor de beeldmodaliteit.
     Dit voorkomt interferentie tussen klassen en modaliteiten tijdens het leren.

2. Ruimtelijke Ankering (Space Anchoring):
   Om te voorkomen dat de leerbare prompts afwijken van klinisch betekenisvolle semantiek, worden deze "verankerd" aan LLM-genereren klinische tekstbeschrijvingen.

   • Een Large Language Model (LLM) genereert voor elke klasse visuele beschrijvingen.
   • Deze beschrijvingen worden omgezet in vaste tekstuele prototypes (anchors).
   • Een verliesfunctie ($L_{anc}$) minimaliseert de afstand tussen de leerbare prompts en deze vaste prototypes, zodat de prompts niet "wegdrijven" naar modality-specifieke ruis.

3. Cross-Modale Manifold-uitlijning via FGW:
   Dit is het meest innovatieve onderdeel. Om de structurele relaties van de hoogwaardige ruimte over te dragen naar de laagwaardige ruimte, gebruiken de auteurs Fused Gromov-Wasserstein (FGW) optimal transport.

   • De hoogwaardige tekstuele embeddings fungeren als een vaste referentiemanifold.
   • FGW dwingt de laagwaardige prompt-manifold om de relatieve geometrische structuur van de hoogwaardige ruimte na te bootsen.
   • Dit zorgt ervoor dat de beslissingsgrenzen in de laagwaardige ruimte consistent blijven met de klinische logica van de hoogwaardige ruimte, zonder dat er visuele data van de laagwaardige modaliteit nodig is.

De totale objectieve functie combineert cross-entropy verlies (voor de hoogwaardige data), anker-verlies en de FGW-uitlijningsverlies.

Belangrijkste Bijdragen

1. Strikte Zero-Shot Strategie: Een asymmetrische overdrachtsmethode die uitsluitend vertrouwen heeft op hoogwaardige visuele data en LLM-genereren tekst, volledig eliminerend de noodzaak voor trainingsdata van de doel-modaliteit.
2. Preventie van Catastrofaal Vergeten: Door middel van semantische ankering en structurele uitlijning wordt voorkomen dat het model overfit op bron-domein statistieken.
3. FGW voor Manifold-uitlijning: Een nieuwe doelstelling die de relationele geometrie van prompts tussen modaliteiten uitlijnt, wat een effectieve route biedt voor zero-shot cross-modale deploy.
4. Factorisatie van Prompts: Het introduceren van gescheiden context-vectoren voor klassen en modaliteiten om beter te kunnen generaliseren.

Resultaten

Het framework is geëvalueerd op vier diverse cross-modale benchmarks:

• Dermoscopie $\to$ Klinische foto's (Huidlaesies)
• Mammografie $\to$ Echografie (Borstlaesies)
• CT $\to$ Borst X-ray (COVID-19 pneumonie)

Kernresultaten:

• State-of-the-art (SOTA) Prestaties: K-MaT behaalde een gemiddelde harmonische mean (H) van 44,1% voor nauwkeurigheid en 36,2% voor macro-F1, wat een verbetering is ten opzichte van BiomedCoOp (respectievelijk 42,0% en 35,0%).
• Preventie van Vergeten: Op de uitdagende borstbeeldvormingstaak (Mammografie $\to$ Echografie) daalde de nauwkeurigheid van standaardmethoden zoals CoOp van 75,2% (bron) naar slechts 27,0% (doel). K-MaT hield de doel-nauwkeurigheid op 38,4% en behaalde een harmonische mean van 50,3%.
• Ablatie-studies: De studie toonde aan dat zowel de anker-methode ($L_{anc}$) als de FGW-uitlijning ($L_{fgw}$) essentieel zijn. Zonder FGW faalt het model om de structurele relaties over te dragen, wat leidt tot een instorting van de prestaties op de doel-modaliteit.

Betekenis en Conclusie

K-MaT biedt een krachtige oplossing voor het probleem van domain shift in medische AI, specifiek wanneer er geen gelabelde data beschikbaar is voor de doel-modaliteit (bijv. in lagere zorginstellingen of ontwikkelingslanden). Door de beslissingsstructuur van rijke, hoogwaardige data (zoals CT) te "transporteren" naar armere modaliteiten (zoals X-ray) via tekstuele kennis en optimal transport, maakt het framework robuuste zero-shot diagnose mogelijk.

Hoewel de absolute prestaties op laagwaardige modaliteiten nog beperkt zijn vergeleken met de bron-modaliteit, is de verbetering in generalisatie significant. De auteurs wijzen erop dat toekomstig werk gericht moet zijn op het integreren van betrouwbaardere visuele signalen om de kloof tussen tekstuele ankers en visuele realiteit in extreme modality-verschillen verder te overbruggen.