Unified Medical Image Tokenizer for Autoregressive Synthesis and Understanding

Probleemstelling

De snelle evolutie van autoregressieve (AR) modellen (zoals GPT-4o) heeft multimodaal AI revolutionair gemaakt. Echter, de toepassing hiervan op medische beeldvorming wordt beperkt door het ontbreken van een unificerende image tokenizer die tegelijkertijd:

1. Fijne anatomische structuren behoudt (noodzakelijk voor beeldreconstructie en synthese).
2. Rijke klinische semantiek encodeert (noodzakelijk voor diagnose en begrip).

Bestaande benaderingen vallen in twee categorieën die beide tekortschieten voor medische doeleinden:

• Synthese-georiënteerde tokenizers (bijv. VQGAN): Deze zijn uitstekend in pixel-niveau reconstructie maar missen discriminatieve, hoog-niveau concepten.
• Begrips-georiënteerde tokenizers (bijv. CLIP): Deze zijn goed in tekstuele semantiek maar verliezen vaak ruimtelijke structuur en textuurdetails.

Daarnaast zijn bestaande methoden voor het trainen van een unificerende tokenizer vaak gebaseerd op gezamenlijke optimalisatie van reconstructie en semantische doelen. Dit leidt tot gradient interference (onderlinge storing) en suboptimale prestaties. Bovendien is er in de medische domein een groot tekort aan gepaarde data (afbeelding + beschrijving) vergeleken met de overvloed aan ongepaarde medische afbeeldingen, waardoor bestaande methoden die grote hoeveelheden gepaarde data vereisen, niet schaalbaar zijn.

Methodologie: MedITok

De auteurs stellen MedITok voor, de eerste unificerende medische image tokenizer, en introduceren een principiële twee-staps trainingsframework om de bovengenoemde problemen op te lossen. Het idee is om visuele representatie als brug te gebruiken.

1. Architectuur:
MedITok bestaat uit een beeldencoder ($E$), een quantizer ($Q$) en een decoder ($D$). Het gebruikt Vector Quantization (VQ) om afbeeldingen om te zetten in discrete token-sequenties, geschikt voor autoregressieve modellen.

2. Het Twee-Staps Trainingsframework:
In plaats van alle doelen tegelijk te trainen, wordt het proces opgesplitst:

• Fase 1: Visuele Representatie Alignering (Visual Representation Alignment)

  • Doel: Het opbouwen van een robuuste codering van laag-niveau structurele informatie en het waarborgen van hoge reconstructie-echtheid.
  • Data: Gebruikt een enorme corpus van ongepaarde medische afbeeldingen (33,4 miljoen).
  • Methode: De tokenizer wordt getraind op reconstructie, maar met een lichte semantische constraint geleid door een vooraf getrainde visuele encoder (BiomedCLIP-V). Dit zorgt ervoor dat de latent space fundamentele semantiek aanneemt zonder de reconstructie te verstoren.
  • Voordeel: Hierdoor wordt de overvloed aan ongepaarde medische data effectief benut.

• Fase 2: Tekstuele Semantische Alignering (Textual Semantic Alignment)

  • Doel: Het verfijnen van de tokens met fijne, hoog-niveau klinische semantiek.
  • Data: Gebruikt gepaarde afbeelding-tekst paren (2,4 miljoen).
  • Methode: De tokenizer wordt verder getuned door de tokens uit te lijnen met fijne-granulariteit tekstuele embeddings van klinische rapporten (via een vooraf getrainde tekstencoder, BiomedCLIP-T).
  • Voordeel: Dit injecteert specifieke medische kennis in de latent space zonder de reeds opgebouwde structurele integriteit te vernietigen.

3. Dataverzameling:
De auteurs hebben een schaalbaar dataset samengesteld met meer dan 33 miljoen medische afbeeldingen en 2,4 miljoen afbeelding-tekst paren, dekkend 9 verschillende beeldvormingsmodaliteiten (zoals CT, MRI, X-ray, dermoscopie, endoscopie, etc.) en diverse anatomische regio's.

Belangrijkste Bijdragen

1. Nieuw Trainingsframework: Een innovatieve twee-staps aanpak die de conflict tussen reconstructie en semantisch leren oplost en het gebruik van ongepaarde data mogelijk maakt.
2. MedITok: De eerste unificerende tokenizer voor medische beelden die zowel laag-niveau details als hoog-niveau klinische semantiek in één latent space encodeert.
3. Uitgebreide Evaluatie: Prestaties zijn getest op meer dan 30 benchmarks, 9 modaliteiten en 4 taakfamilies (reconstructie, classificatie, synthese, VQA).
4. Open Source: Het model en de code zijn open-source beschikbaar gesteld.

Resultaten

MedITok presteert state-of-the-art (SOTA) op diverse taken:

• Beeldreconstructie (Low-Level): MedITok behaalt de beste resultaten in PSNR, SSIM en rFID over 8 modaliteiten, zelfs bij een downsampling factor van 16x. Het behoudt kleuren en fijne structuren beter dan gespecialiseerde medische tokenizers (zoals MedVAE) of algemene tokenizers.
• Diagnostische Informatiebehoud: Via een proxy-classificatietest (ResNet op gereconstrueerde beelden) toont MedITok aan dat het diagnostisch cruciale details behoudt, met significante verbeteringen in mAP en AUC ten opzichte van concurrenten.
• Semantische Codering (High-Level): Bij lineaire probing voor beeldclassificatie (bijv. huidlaesies, longontsteking) overtreft MedITok alle andere tokenizers, wat aantoont dat de latent space rijk is aan klinisch relevante semantiek.
• Medische Beeldsynthese: Geïntegreerd in een autoregressief model (LlamaGen), genereert MedITok realistischere en diverse medische beelden met een lagere gFID (Fréchet Inception Distance) dan modellen gebaseerd op VQGAN of UniTok.
• Medisch Begrip (VQA): In Visual Question Answering taken (geïntegreerd in LLaVA-Med) behaalt MedITok de hoogste nauwkeurigheid op benchmarks zoals VQA-RAD en SLAKE, wat aantoont dat het model klinische concepten correct kan interpreteren.

Betekenis en Impact

Deze werken markeert een belangrijke stap in de ontwikkeling van foundation models voor de medische sector.

• Democratisering: Het biedt een schaalbare basis voor het bouwen van krachtige autoregressieve modellen voor medische toepassingen, zonder dat elke ontwikkelaar een eigen tokenizer hoeft te trainen.
• Unificatie: Het overbrugt de kloof tussen beeldgeneratie (synthese) en beeldbegrip (diagnose) in één enkel framework.
• Efficiëntie: Door het slimme gebruik van ongepaarde data in Fase 1, maakt het trainen van dergelijke modellen haalbaar in een domein waar gelabelde data schaars is.

Kortom, MedITok fungeert als een fundamentele bouwsteen voor de volgende generatie multimodale AI-systemen in de gezondheidszorg, met de potentie om workflows te stroomlijnen en patiëntuitkomsten te verbeteren door zowel realistische synthese als nauwkeurige diagnoseondersteuning mogelijk te maken.