⚡ electrical engineering

Unified Medical Image Tokenizer for Autoregressive Synthesis and Understanding

This paper introduces MedITok, a unified medical image tokenizer trained on over 33 million images across nine modalities using a novel two-stage framework that leverages unpaired data for reconstruction fidelity before injecting textual semantics, thereby enabling state-of-the-art autoregressive synthesis and understanding for diverse diagnostic and generative applications.

Original authors: Chenglong Ma, Yuanfeng Ji, Jin Ye, Zilong Li, Chenhui Wang, Junzhi Ning, Wei Li, Lihao Liu, Qiushan Guo, Tianbin Li, Junjun He, Hongming Shan

Published 2026-04-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Chenglong Ma, Yuanfeng Ji, Jin Ye, Zilong Li, Chenhui Wang, Junzhi Ning, Wei Li, Lihao Liu, Qiushan Guo, Tianbin Li, Junjun He, Hongming Shan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a super-smart AI robot how to be a doctor. This robot needs to do two very different things:

Look at an X-ray or MRI and describe exactly what it sees (like a radiologist saying, "I see a shadow here that looks like pneumonia").
Look at a description and draw a brand-new, realistic medical image (like an artist sketching a lung based on a doctor's notes).

The problem is that the robot's "eyes" (the part that turns pictures into data the AI can understand) have been broken up into two separate tools. One tool is great at drawing pictures but terrible at understanding medical jargon. The other is great at understanding words but terrible at drawing the fine details of a bone or a tumor.

This paper introduces MedITok, a new "universal translator" for medical images that fixes this problem. Here is how it works, explained simply:

The Problem: The "Split Personality" Translator

Think of existing image tools like a translator who only speaks two languages but not at the same time:

The Artist: Can copy a picture perfectly pixel-by-pixel but doesn't know what a "fracture" or "inflammation" means.
The Scholar: Can write a perfect essay about a disease but can't draw a picture of it to save their life.

If you try to teach the robot to be both an Artist and a Scholar at the exact same time, they argue. The Artist wants to focus on the tiny lines of a bone, while the Scholar wants to focus on the big concept of "disease." They get in each other's way, and the robot learns poorly. Also, in the medical world, we have millions of X-rays without descriptions, but very few X-rays with descriptions. Trying to learn from just the few descriptions isn't enough.

The Solution: A Two-Step Training Camp

The authors created a clever two-step training camp for MedITok so it can learn to be both an Artist and a Scholar without them fighting.

Step 1: The "Visual Gym" (Using Millions of Unlabeled Photos)
First, they take the robot to a gym with 33 million medical images. None of these images have text descriptions attached.

The Goal: The robot just has to look at the image and try to rebuild it perfectly from memory.
The Trick: While it's doing this, they give it a "hint" from a pre-trained expert (a visual AI that already knows what things look like). This ensures the robot learns the structure: the shape of a heart, the texture of a bone, the color of a skin lesion.
Analogy: Imagine an art student practicing by copying thousands of paintings without reading the titles. They learn how to hold the brush, mix the colors, and capture the light. They become a master painter.

Step 2: The "Medical School" (Using Paired Images and Text)
Now that the robot is a master painter, they bring in the 2 million images that do have descriptions (captions).

The Goal: The robot now learns to connect the pictures it already knows how to draw with the specific words doctors use.
The Trick: They align the robot's "painting" with the "text." When the robot sees a picture of a cloudy lung, it learns that the specific pattern of pixels corresponds to the words "ground-glass opacity."
Analogy: Now the art student goes to medical school. They take their perfect painting skills and learn the vocabulary. They learn that this specific shade of gray means "pneumonia."

The Result: The "MedITok" Super-Tool

By separating these steps, MedITok becomes a Unified Medical Image Tokenizer.

Token: Think of this as a LEGO brick. MedITok turns a whole medical image into a sequence of LEGO bricks.
The Magic: These bricks contain both the fine details (so you can rebuild the image perfectly) and the medical meaning (so the AI can understand the diagnosis).

Why This Matters

Because MedITok is so good at this, it can power the next generation of medical AI:

Better Diagnosis: It can help AI read X-rays and explain why it thinks a patient is sick, with high accuracy.
Synthetic Data: It can generate fake but realistic medical images. Why is this useful? Doctors can use these fake images to train new AI models without needing to find millions of real patient records, which protects patient privacy.
Universal: It works on 9 different types of medical scans (X-rays, MRIs, CTs, ultrasounds, etc.), making it a versatile tool for the whole medical field.

In a Nutshell

The authors built a bridge. Instead of forcing an AI to learn to draw and speak at the same time (which causes confusion), they let it learn to draw first using a massive library of pictures, and then taught it to speak using a smaller library of picture-text pairs. The result is a single, powerful tool that can both see and understand medical images better than anything before it.

1. Problem Statement

The rapid advancement of autoregressive (AR) models in multimodal AI (e.g., GPT-4o) has demonstrated the potential for unified image synthesis and understanding. However, applying these models to medical imaging faces two critical bottlenecks:

Lack of Unified Tokenizers: Existing tokenizers are specialized. Synthesis-oriented tokenizers (e.g., VQGAN) preserve low-level structural details but lack high-level clinical semantics. Understanding-oriented tokenizers (e.g., CLIP) capture semantic concepts but fail to retain fine-grained spatial structures and textures required for medical diagnosis.
Data Scarcity and Interference: Training a unified tokenizer typically requires joint optimization of image reconstruction and textual semantic objectives. In the medical domain, high-quality paired image-text data is scarce compared to abundant unpaired images. Furthermore, joint optimization often leads to gradient interference, where the competing objectives (pixel fidelity vs. semantic alignment) result in suboptimal performance or representation collapse.

2. Methodology: MedITok and Two-Stage Training

The authors propose MedITok, the first unified medical image tokenizer, trained via a novel two-stage framework designed to decouple structural learning from semantic learning while leveraging both unpaired and paired data.

Architecture

MedITok follows a Vector Quantization (VQ) based architecture consisting of:

Image Encoder ( $E$ ): Maps input images to continuous latent vectors.
Quantizer ( $Q$ ): Discretizes latent vectors into a finite codebook (32,768 tokens).
Image Decoder ( $D$ ): Reconstructs images from quantized tokens.
Pretrained Encoders: Uses BiomedCLIP (frozen) as the semantic bridge, providing both vision ( $E_{vis}$ ) and text ( $E_{txt}$ ) encoders.

The Two-Stage Training Framework

Instead of joint optimization, the framework uses a curriculum learning approach:

Stage 1: Visual Representation Alignment (Unpaired Data)

Goal: Establish a robust latent space that preserves low-level structural fidelity while introducing light semantic constraints.
Data: Trained on 33.4 million unpaired medical images across 9 modalities (CT, MRI, X-ray, Dermoscopy, etc.).
Mechanism: Optimizes for image reconstruction ( $L_{rec}$ ) while adding a contrastive loss ( $L_{vis}$ ) that aligns the latent tokens with the features of a pretrained vision encoder ( $E_{vis}$ ).
Outcome: Creates a high-fidelity codebook capable of accurate reconstruction, serving as a "bridge" for semantic injection.

Stage 2: Textual Semantic Alignment (Paired Data)

Goal: Inject fine-grained clinical semantics into the established latent space.
Data: Trained on 2.4 million image-text pairs (captions/reports).
Mechanism: Refines the encoder using a contrastive loss ( $L_{txt}$ ) that aligns image tokens with detailed clinical captions processed by a pretrained text encoder ( $E_{txt}$ ).
Outcome: The tokenizer now encodes both structural details and rich clinical semantics, enabling downstream AR models to perform both generation and comprehension.

3. Key Contributions

Novel Training Framework: Introduced a principled two-stage training strategy that uses visual representation alignment as a bridge. This effectively mitigates the gradient interference between reconstruction and semantic objectives and allows the utilization of massive unpaired medical datasets.
MedITok Model: The first unified medical image tokenizer that successfully encodes both low-level structural details (for synthesis) and high-level clinical semantics (for understanding) within a single discrete latent space.
Large-Scale Dataset Curation: Curated a massive, diverse dataset comprising 33.4 million medical images and 2.4 million image-text pairs spanning 9 imaging modalities and diverse anatomical regions, with rigorous quality control to filter out non-clinical content.
Open Source: The model, code, and data access protocols are made available to the community to democratize medical foundation model development.

4. Experimental Results

MedITok was evaluated on 30+ benchmarks across 9 imaging modalities and 4 task families, outperforming state-of-the-art general-domain tokenizers (UniTok, TokenFlow) and specialized medical tokenizers (MedVAE, PUMIT).

Low-Level Encoding (Reconstruction):
- Achieved the best rFID, PSNR, and SSIM across 8 modalities (including CT, MRI, X-ray, and Dermoscopy).
- Preserved diagnostic information better than baselines, evidenced by higher mAP and AUC in proxy classification tasks on reconstructed images (e.g., skin lesion diagnosis, pneumonia detection).
High-Level Encoding (Classification):
- In linear probing tasks, MedITok achieved the highest classification accuracy (mAP/AUC) across dermoscopy, fundus, pathology, ultrasound, and X-ray, demonstrating superior semantic representation.
Medical Image Synthesis:
- Integrated into LlamaGen (an AR image generator), MedITok produced images with the lowest gFID (best fidelity and diversity).
- A "Visual Turing Test" with radiologists showed that MedITok-generated images were harder to distinguish from real images compared to those generated by UniTok.
Medical Image Understanding (VQA):
- When integrated into LLaVA-Med, MedITok achieved state-of-the-art accuracy on VQA-RAD and SLAKE benchmarks.
- It demonstrated superior ability to ground clinical terms to visual features (e.g., correctly identifying "bilateral pleural effusions" vs. "right-sided" or hallucinating colors).

5. Significance and Impact

Foundation for Medical AR Models: MedITok provides a scalable, unified component for next-generation multimodal medical AI, enabling models that can simultaneously synthesize realistic medical images and interpret complex clinical reports.
Solving Data Imbalance: The two-stage framework offers a practical solution to the scarcity of paired medical data by effectively leveraging the abundance of unpaired images, a critical bottleneck in medical AI.
Democratization: By open-sourcing a high-performance tokenizer trained on diverse modalities, the work lowers the barrier for researchers to build specialized medical AR models without needing to train tokenizers from scratch.
Future Directions: The authors note that while MedITok is currently 2D-focused, the framework is extensible to 3D volumetric data and can be adapted for more complex downstream workflows, paving the way for general-purpose medical AI systems.