Align Your Query: Representation Alignment for Multimodality Medical Object Detection

Probleemstelling

Medische objectdetectie (het lokaliseren en classificeren van anatomische structuren of afwijkingen) lijdt onder prestatieverlies wanneer een enkel model wordt getraind op een gemengde dataset van verschillende beeldvormingsmodaliteiten (bijv. röntgenfoto's (CXR), CT-scans, MRI, endoscopie, pathologie).

• Heterogeniteit: Elke modaliteit heeft unieke statistische eigenschappen en visuele kenmerken, wat leidt tot disjuncte representatieruimtes.
• Prestatiedaling: Een model dat op één modaliteit is getraind, faalt vaak bij generalisatie naar andere modaliteiten. Het trainen van één model op een gemengde dataset resulteert vaak in een verward representatieruimte waar het model moeite heeft om modaliteit-specifieke context te onderscheiden, wat de detectienauwkeurigheid verlaagt.

Methodologie

De auteurs stellen een detector-onafhankelijk framework voor dat zich richt op representatie-alignment (uitlijning van representaties) van de object queries in moderne transformer-gebaseerde detectoren (zoals DETR-varianten). Het framework bestaat uit twee hoofdcomponenten:

1. Modality Tokens (Modaal Tokens)

In plaats van complexe architecturale wijzigingen, worden compacte, tekst-gebaseerde embeddings gebruikt om de modaliteit en de klasse expliciet te coderen.

• Constructie: Voor elke combinatie van modaliteit ($d$) en klasse ($c$) wordt een tekst-prompt gegenereerd (bijv. "Aortic enlargement in CXR").
• Embedding: Een vooraf getrainde tekst-encoder (zoals CLIP, BiomedCLIP of PubMedCLIP) verwerkt deze tekst tot een embedding. De [CLS]-token wordt gebruikt en via een lineaire projectie omgezet in een modality token ($m_{d,c}$).
• Voordeel: Deze tokens zijn lichtgewicht, vereisen geen extra annotaties en fungeren als stabiele ankers voor de modaliteit.

2. Multimodality Context Attention (MoCA)

Om de modale context te integreren in het detectieproces zonder de architectuur van de detector te veranderen, wordt MoCA voorgesteld.

• Mechanisme: De modality token wordt direct aan de set van object-queries van de detector toegevoegd.
• Self-Attention: In de decoder van de detector (bijv. DINO) wordt een self-attention-laag gebruikt waarbij de queries en de modality token samen worden verwerkt. Hierdoor kunnen de object-queries informatie "attenderen" op de modality token.
• Resultaat: Dit injecteert modaliteit-specifieke context direct in de queries op het moment dat de beslissingen worden genomen, met verwaarloosbare vertraging (latency).

3. QueryREPA (Query Representation Alignment)

Om de uitlijning te versterken, wordt een pre-trainfase geïntroduceerd voordat de detectie-head wordt getraind.

• Doel: Het uitlijnen van de gemiddelde query-representatie van een afbeelding met de bijbehorende modality token.
• Verliesfunctie: Er wordt een contrastief verlies (InfoNCE) gebruikt. De gemiddelde query wordt geprojecteerd naar de ruimte van de modality tokens en geoptimaliseerd om dicht bij de juiste token te liggen en ver weg van tokens van andere modaliteiten.
• Modality Batch Sampling: Om te voorkomen dat één modaliteit de training domineert, worden mini-batches samengesteld die een gebalanceerde mix van verschillende modaliteiten bevatten. Dit zorgt voor sterke contrasten tussen modaliteiten.
• Fase: Dit gebeurt als een korte pre-trainfase op (gedeeltelijk) ongelabelde data, waarna de volledige detector wordt gefinetuned.

Belangrijkste Bijdragen

1. Modality Tokens: Een eenvoudige, detector-onafhankelijke methode om modaliteitscontext expliciet te maken met minimale overhead en zonder extra annotaties.
2. MoCA: Een "drop-in" self-attention mechanisme dat modality tokens aan queries toevoegt en deze mixt binnen de decoder, waardoor multimodale detectie wordt versterkt zonder architecturale wijzigingen.
3. QueryREPA: Een pre-trainstrategie die query-representaties uitlijnt met modality tokens via een contrastief doel. Dit maakt gebruik van grote, ongelabelde datasets en verbetert de downstream prestaties aanzienlijk.
4. Synergie: Het paper toont aan dat MoCA en QueryREPA complementair zijn en samen een sterke synergie-effect hebben.

Resultaten

De methode is getest op een geaggregeerde dataset met diverse modaliteiten (CXR, CT, MRI, Colonoscopie, Pathologie) en vergeleken met state-of-the-art detectoren (Deformable DETR, Sparse R-CNN, DiffusionDet, GLIP, Grounding DINO, DINO).

• Prestatieverbetering: De voorgestelde methode (DINO + MoCA + QueryREPA) bereikte een nieuwe state-of-the-art (SOTA) op de gemengde dataset.
  • De totale AP (Average Precision) steeg van 37.7 (baseline DINO) naar 41.3 (met PubMedCLIP tokens).
  • De AP50 steeg van 58.6 naar 65.5.
• Robuustheid: De verbetering was consistent ongeacht de gebruikte tekst-encoder (CLIP, BiomedCLIP, PubMedCLIP) of de basisdetector (Deformable DETR, DINO).
• Efficiëntie: De inferentie-tijd nam slechts marginaal toe (bijv. van 9.37ms naar 9.59ms voor DINO), wat aantoont dat de methode zeer efficiënt is.
• Kwalitatieve Resultaten: Het model slaagde er beter in om subtiele laesies te detecteren en nauwkeurigere bounding boxes te genereren in vergelijking met baselines, vooral bij kleine objecten en complexe achtergronden.
• Gebruik van Ongelabelde Data: QueryREPA bleek effectief bij het gebruik van datasets met alleen klasse-labels (zonder bounding boxes) voor pre-training, wat een nieuwe weg opent voor zwak toezicht in de medische beeldanalyse.

Betekenis en Conclusie

Dit paper biedt een praktische en schaalbare oplossing voor het probleem van heterogeniteit in medische beeldvorming. Door in te zetten op de uitlijning van object queries in plaats van het opnieuw trainen van zware backbones of het gebruik van complexe cross-attention mechanismen, maken de auteurs het mogelijk om robuuste, generaliseerbare detectiemodellen te bouwen voor diverse medische modaliteiten.

De aanpak is bijzonder relevant voor de klinische praktijk, waar data vaak versnipperd is over verschillende modaliteiten en waar annotaties schaars zijn. De methode opent de deur naar "generalist" medische AI-modellen die betrouwbaar kunnen werken in gemengde omgevingen zonder de architectuur van bestaande detectoren fundamenteel te hoeven veranderen.